Flexible Docking via Unbalanced Flow Matching

Diffusion models have emerged as a recent successful paradigm for molecular docking. However, these methods treat the protein either as a rigid structure, or force the model to fold proteins from unstructured noise. In this work, we instead focus on flexible docking, leveraging the unbound distribution of proteins to model the precise effect(s) of ligand binding. While Flow Matching (FM) presents an attractive option for this task, we show that a naive application of flow matching results in a complex learning task with poor performance. We thus propose Unbalanced Flow Matching, a generalization of flow matching that allows us to tradeoff sample efficiency with approximation accuracy by relaxing the marginal constraints. Empirically, we validate our framework on flexible docking, demonstrating strong improvements in protein conformation prediction while retaining comparable docking accuracy.