Composing Unbalanced Flows for Flexible Docking and Relaxation
Authors: Gabriele Corso
Vignesh Ram Somnath
Noah Getz
Regina Barzilay
Tommi Jaakkola
Andreas Krause
Venue: ICLR 2025 Oral (2025)
Keywords: molecular docking
flow matching
structure relaxation
unbalanced transport
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: A new generalized flow matching paradigm and its applications to flexible docking and relaxation

Abstract: Diffusion models have emerged as a successful approach for molecular docking, but they often cannot model protein flexibility or generate nonphysical poses. We argue that both these challenges can be tackled by framing the problem as a transport between distributions. Still, existing paradigms lack the flexibility to define effective maps between such complex distributions. To address this limitation, we propose Unbalanced Flow Matching, a generalization of Flow Matching (FM) that allows trading off sample efficiency with approximation accuracy and enables more accurate transport. Empirically, we apply Unbalanced FM on flexible docking and structure relaxation, demonstrating our ability to model protein flexibility and generate energetically favorable poses. On the PDBBind docking benchmark, our method FlexDock improves the docking performance while increasing the proportion of energetically favorable poses from 30% to 73%.

Published as a conference paper at ICLR 2025
COMPOSING UNBALANCED FLOWS FOR FLEXIBLE DOCKING AND RELAXATION
Gabriele Corso MIT
Vignesh Ram Somnath ETH Zurich
Noah Getz MIT
Regina Barzilay MIT
Tommi Jaakkola MIT
Andreas Krause
ETH Zurich
ABSTRACT
Diffusion models have emerged as a successful approach for molecular docking, but they often cannot model protein flexibility generate nonphysical poses. We argue that both these challenges can be tackled by framing the problem as a trans- port between distributions. Still, existing paradigms lack the flexibility define effective maps between such complex distributions. To address this limitation, we propose Unbalanced Flow Matching, generalization of Flow Matching (FM) that allows trading off sample efficiency with approximation accuracy and enables more accurate transport. Empirically. we apply Unbalanced FM on flexible docking and structure relaxation, demonstrating our ability to model protein flexibility and generate energetically favorable pose On the PDBBind docking benchmark, our method FLEXDOCK improves the docking performance while increasing the proportion of energetically favorable poses from 30% to 73%.
INTRODUCTION
The mechanism of action of most drug molecules is often rationalized by studying how the molecule binds to one or more key proteins. Computationally, predicting the structure of these binding interactions constitutes the molecular docking task. Over the past decades, significant progress has been made in molecular docking, initially through classical search techniques and more recently with deep learning models. However, these methods primarily focus on rigid docking, assuming the protein has a fixed structure.
Proteins however are flexible entities and often undergo conformational changes during docking. Computational methods for understanding flexibility fall into two categories: co-folding and flexible docking. Co-folding involves predicting the bound structure of the protein and the ligand from scratch as a single task. Flexible docking, instead, aims at only modeling the limited structural transformation between unbound and bound protein structures, formulation that allows for more efficient, interpretable, and controllable methods.
However, existing flexible docking methods have far failed to provide satisfactory levels of accuracy. Among these, search-based techniques struggle to account for protein degrees of freedom due to the significantly increased dimensionality of the search space (Koes al. 2013 McNutt et al. 202T). Deep learning methods have improved on this by extending diffusion or flow processes to include the protein flexibility, but these either severely limit the degrees-of-freedd force the model to learn protein folding entirely, resulting in structure predictions that are frequently inaccurate (Qia et al.) 2024 Lu al. 2024). Furthermore, recent works have highlighted the tendency of ML docking models to generate non-physical structures (Buttenschoen et al. 2024).
implicitly
Addressing the challenges above requires new techniques for precise transport between the complex distributions of unbound (apo) and bound (holo) protein structures. Flow matching (FM) (Lipman al. 2022 Albergo al. 2023). in principle, can learn such a transport. However, we show that
"Equal contribution. Correspondence to wco. wtp.ii.mit..cm edu and vsomnath@inf www.iii..eeethhhhh.eeo OOr code and models are available httt:p///gggttbb...o//////mata//

---

Published as a conference paper at ICLR 2025
Manifold Docking
Structure Relaxation
Momioll Unbalanced Flow Matching AA
Euclidean Unbelennzed Matching: 24
precicted structurcs
DO struchieee
Figure I: Schematic representation of the FLEXDOCK docking framework formed by the chaining of two unbalanced flows. The first (details in Section [4.1) is aimed resolving the approximate pose by working on longer flows but reduced degrees of freedom, and the second (Section 4.2) aims at relaxing the generated structure producing high-quality poses.
directly applying flow matching this task results in complex mappings, difficult learning task, and consequently poor performance. To address this, we propose Unbalanced Flow Matching, where relax the marginal constraints of FM through more careful choice of coupling distribution. We study these choices both formally and empirically, resulting in a higher docking performance, We further consider composing short unbalanced flows and show that, in the limit, this corresponds to local likelihood gradient steps towards the desired marginal.
Based on the insights above, we construct new flexible pocket-based docking method FLEXDOCK by chaining two separate Unbalanced Flows (see Figure The first, referred to as manifold docking, defined on low dimensional space and aims to predict the approximate ligand pose and protein structural change. The second, relaxation, is defined on the full Euclidean space but with shorter maps and aims to provide a fast yet effective relaxation of the structure.
Empirically, we demonstrate that the new modeling perspective of FLEXDOCK enhances structure prediction quality, especially for protein conformations. On the PDBBind flexible pocket-based dock- ing benchmark, FLEXDOCK improves the proportion of very accurate protein structure predictions (all-atom RMSD below IÅ) from 32% of DiffDock-Pocker to 42% while having the same training data and model architecture. Furthermore, thanks the relaxation flow, the generated poses are far more physically realistic, as demonstrated by the proportion of poses passing PoseBusters checks (Buttenschoen et al. 2024) which increases from 30% to 73%.
In summary, propose Unbalanced Flow Matching, new distribution matching framework, and prove provides a way to balance sample efficiency and sample quality. Furthermore, by chaining together new manifold docking and structure relaxation models based on Unbalanced FM, we develop a flexible docking method that addresses two major limitations previous docking techniques: accurate protein flexibility modeling and the generation of energetically favorable poses.
BACKGROUND AND RELATED WORK
Flow matching. Given two distributions go and Q1, Flow Matching (Lipman et al. 2022 Albergo
et al. 2023) provides way of learning vector field v, which induces a continuous normalizing flow
(x) that transports qu to Q1, i.e., qı (x) [v1]+ (x)), where # denotes the pushforward operator.
The key idea in FM is to define conditional flow 1h (xooo))) interpolating between Xo 9o and
X1 91. and its associated vector field u, (xxxx,) ) & the (xx One can then learn the marginal vector field (x) with neural network Vi (x; f), by regressing against the conditional vector field with the conditional flow matching (CFM) objective:
LCFM Er [x;; 0) ux(xx(/xx) (x1]]
FM was further generalized by Pooladian et al. (2023) and Tong et al. (202)) who show that the CFM objective does not require independent samples from go and g1 but it can use an arbitrary joint sampling distribution (xx,, x1), the coupling distribution, long it satisfies the marginal
2

---

Published as a conference paper at ICLR 2025
constraints being qo and @1 respectively. This formulation enables drawing connection between FM and optimal transport (OT). When using OT to define the coupling distribution q. the flows become
straight and the transport cost (x,)) (xo) - Xo l is the OT cost W3 go. 91). An extension of flow matching to the unbalanced setting was first explored in (Eyring et al. 2024), with couplings q
estimated using Unbalanced OT(UOT). In this work, we focus on arbitrary couplings and provide a formal and empirical characterization of the associated objectives.
Flexible Docking, Flexible docking assumes access to unbound structures of proteins (known as apo) and predicts how these will change upon ligand binding (producing holo-structures). Traditional search-based docking methods define a scoring function and search the space of possible poses, aiming to find the pose minimizing the scoring function (Alhossary et al. 2013 McNutt al.) 202I). These methods can in principle incorporate protein flexibility by adding torsion angles of the sidechains in the pocket to the search space. However, they struggle to find optimal joint poses due to the increased dimensionality and they are unable to model proteins' flexibility beyond the sidechains.
Recently, a few deep learning methods flexible docking have been proposed. DIFF -PCCKT (Plainer et al. 2023) and RF-DocK (Huang et al., (024) use diffusion and diffusion bridge models to capture the ficxibility in protein pocket sidechains addition to ligand flexibility, they, however, once again cannot model proteins' flexibility beyond sidechain torsion angles. DYNAMICBIND (Lu et al., 2024) incorporates backbone flexibility but uses hardcoded noise perturbation rules to interpolate from apo residue frames to holo residue frames. [Somnath al. (2023) explicitly predicts the conformational changes between apo and holo states of proteins using Diffusion Schrödinger Bridges by treating the apo and holo states paired data.
Structure Relaxation. Even in rigid docking, ML methods, despite improving geometric prediction accuracy, often produce poses with nonphysical properties like sterie clashes and strained bonds (Buttenschoen et al.] 2024). These artifacts can hinder downstream prediction tasks, necessitating relaxation using molecular dynamics force fields. However, this relaxation process is often long (Bryant et al. 2024). complicating its use in large-scale screening.
UNBALANCED FLOW MATCHING
The motivation for flexible docking over co-folding is to leverage the unbound distribution of proteins and focus on modeling the precise effects of ligand binding. Formally, we wish to approximate flow the (equivalently vector field vt), whose pushforward transports the distribution over apo-structures go to the distribution over holo-structures 91.
The flexible docking setting presents two challenges that make a direct application of flow matching for this task ineffective (as we will also demonstrate empirically). First, X-ray crystallography data are the main source of bound conformations, resulting in the availability of single sample for holo-structun for most known complexes. This prevents one from using minibatch- OT based flow matching methods (Pooladian et al., (023) and leads to a large expected length of conditional flows. Second, even if one could construct the optimal transport mapping g between 4o and q), the induced expected cost will likely remain very high. In fact, although the different conformational states of the protein do not change significantly upon ligand binding, their relative weights often notably altered. Training FM in this setting requires the model to move the protein between conformations, leading to long and complex conditional and marginal flows, and harder learning task.
We propose the Unbalanced Flow Matching framework to tackle such problems.
3.1 FORMALIZATION
Let 40 be the input distribution, and 91 target distribution, defined over data space X. Unbalanced Flow Matching (UFM) aims to learn a continuous normalizing flow parameterized with a vector field (xx;;; 8) that minimizes the following objective:
LUFu(g, 0) = -a Er, [xoe(xx; 0) we (xel (xo, x x)))
D2 (xo) + D2 (2(0xxl91
(2)
Mass-variation penalty
Conditional flow matching
3

---

Published as a conference paper at ICLR 2025
(a) Flow Matching
(b) Unbalanced FM
40
q
91
q1
Figure 2: Schematic depiction of the difference between FM and Unbalanced FM. The first builds couplings that maintain the marginal distributions regardless of the difference between the distribu- tions and the consequent expected transport cost. The second trades off some distributional coverage to obtain couplings with significantly reduced transport costs.
where is an arbitrary coupling distribution whose marginals we denote by (xg and (xl, and D2 is Rényi Divergence of order In Section 3.2 we show that this objective describes tractable upper bound to the fundamental approximation efficiency trade-off.
Flow matching and its variants choose a coupling distribution such that marginals are conserved. Although this implies zero mass-variation penalty in Equation the conditional flow matching term can suffer significantly from this choice, as might require the learning of complex couplings and consequently poor approximation. In contrast, our framework jointly accounts for the flow matching and mass-vare varistvaiation penalty terms. However, jointly optimizing and to find global minimum is typically intractable. Our approach thus relies finding an approximation q, and using the flow matching objective with the chosen q. From Equation 2 includes the zero function), is easy to see that:
(and assuming the space ez(:;
min LUFM 0) S min a (x0,,(x0xx)) 9.8
[xlu (xo: x )) + D2(uoluxa) D2(6x
For X = Rdd (x(xxx (Xo, x1)) Xo is common choice. Plugging this into Equation gives us:
min LuFM(q, 0) min Ey, (xoxx)) [o ||x] xxol D2(golaxo) D2(0x1l9)) & UOT(20- 91) 9.8
(4)
The upper bound corresponds to the Unbalanced Optimal Transport (UOT) objective with quadratic cost ||x; xoll² that approximate choose g. Fixing q, then learn 0 via SGD with the resulting simplification of Equation
LUrM q) E2, (x..x)) [|oz 0) u. (xel (xo, x1)
(5)
By lifting the mass-conservation requirement of FM, we now simplify the learning task (see Figure for an empirical validation). However, now even with a perfectly trained vector field, the pushforward of go will no longer correspond to q1- Our goal of obtaining samples from 91. therefore, involves reweighting the samples generated by the learned flow. In practice, we approximate this re-weighting
with discriminator akin to the confidence model used in (Corso al. 2022
(see Algorithm
and Section 4.1
Algorithm 1: UNBALANCED FM INFERENCE
Sample prior Xo (8) 40
Use flow H (-Ig, 0) to map Xo x1
Reweight reject with learned discriminator
Return [x, X1 (2) x..,,x, (n)
4

---

Published as a conference paper at ICLR 2025
3.2 EFFICIENCY VS APPROXIMATION TRADE-OFF
In this section, show how the UFM objective arises from optimizing the trade-off between sample efficiency and approximation accuracy. Sample efficiency refers to the number of samples from Yo needed to obtain approximate sample of 4i (passing the discriminator) under the learned flow. These approximate samples of still incur some imprecision due to the flow not being learned perfectly; the approximation accuracy refers to the Wasserstein distance between @i and the distribution of these approximate samples. To study this, we consider modified version of the Unbalanced FM inference procedure (Algorithm[J). The main difference of this procedure (Algorithm
2 is that we perform some filtering not only after the flow but also before it. Algorithm 2 js, thus, at but easier to study.
most as sample efficient as Algorithm |
Algorithm 2: UFM EFFICIENCY LOWER BOUND
Qx Wal 9x,)
Sample prior Xo (8) qu
p (14.. Iearned Fore
dxg (xo ( ()(x
Reweight or reject based on approx.
Use flow d (-Ig. 9) to map xo (5) to x1 (6) ~ dx1
4,
9s maaginal q
- actimel flo
Reweight or reject based on approx. qi (x,, (1)
9x, (25) (375
Dj6s, ley)
Return 111 (1) x1 (2) ...., X1 (*))
90 starting disorbbttonn
91 sagal disebbation
Algorithm 2: High-level inference routine for modi- fied Unbalanced FM routine used to define theoretical lower bound on the efficiency of the original Unbal- anced FM routine of Algorithm
Figure 3: Relationship between the dif- ferent distributions introduced in the the- oretical analysis of the Unbalanced FM efficiency lower bound algorithm.
Formalization. Let (1((:))) q) be the optimal flow from the Unbalanced FM objective with couplings and (1(;;g 0) its approximation we are able to learn. FM guarantees us that fx: = [Wa (-14) Ixo and let dx. (-10) [,1(-19, #). We summarize all the defined distributions and their relationship
in Figure
The definition of Unbalanced FM as method to bridge two distributions qo and a1 leads us to analyze the trade-off between the approximation error when learning the flow, formalized as W2 (fx, (-10), (x,)) and the sample efficiency ESS* (q) that derives from having to perform rejection sampling to bridge the gaps between 4o and (x. and between Ix1 and 91- Simple mappings will result in low approximation errors but potentially lower efficiency, and vice versa.
Approximation error. Using Theorem 1 from Benton al. (223)) can show, under appropriate assumptions that the approximation error for a given coupling distribution 9. W (9x. (-10), 9x. ), is bounded by FM objective:
W2 (110) S I.2. BE.g [llv:(x;; 0) vn (x.) < T,2 Br, (x(: 0) u. (x (6)
where L = explfe Led() and L, is a constant such that uo (x(x () is L,-Lipschitz in The second inequality follows from the convexity of the squared norm.
Sample efficiency. We can measure the sample efficiency that the model has when converting samples of to unbiased samples of Qı via the effective sample size ESS* (q). i.c., the reciprocal of how many samples from go it takes using the ideal flow (11(;99) q) to generate an unbiased sample from Q1- In Proposition (derivation in Appendix [.I)), we demonstrate that for an ideal flow this sample efficiency is bounded by the similarity between go and qi and the respective marginals of q:
Propasition 1. The effective sample size, ESS*, for sampling q1 when having access to samples of go and a perfectly trained flow with coupling distribution q bounded by:
ESS* (q) exp D2(golax.) D2(gx, [q:))
(7)
5

---

Published as a conference paper at ICLR 2025
Trade-off upper bound. The orf between approximation and sampling efficiency in the choice of optimal flow can then be expressed as a joint objective:
B W (-10), 9x.)
log ESS (g) S LUFM(0,
(8)
Approximation error
Sampling efficiency
where the incquality derives from setting a = L 3 and using the derivation above. This motivates
the UFM objective way to optimize upper bound on the joint approximation error and sample (in-) efficiency trade-off.
In practice, since we only have access to one sample for the distribution over bound structures (the crystal structure PDB, typically unique), we cannot define via exact optimal Unbalanced approximate the optimal coupling with the distribution g(xo, x1) oc OT. In our experiments, where ((Xo,X1) is some, task-specific, transport cost, and Cusk is an go ( (X1 ee(x (x)ll((xx))))xx))))))))))))))))) <iu empirically chosen cutoff to balance sample efficiency and mapping complexity.
3.3
COMPOSING UNRALANCED FLOWS
In Section 4.2 we will show how we obtain effective structure relaxations by learning a further unbalanced flow with smaller matching costs on top of the distribution learned by the initial manifold docking unbalanced flow. In this section, we analyze the effect of composing flows together in the case of short distance couplings that corresponds the setting of structure relaxation.
Each flow tries to learn the matching from the output distribution of the previous flow and the final distribution. We consider the specific coupling introduced in the previous section (Xx., xi) = (x (x, 11 ||S with go now representing the output distribution of the previous flow. In Proposition (derivation in Appendix A..2). we demonstrate that taking e - 0 leads to approximate local likelihood gradient steps that push qo towards matching 91:
2
Proposition 2. A coupling 4(xo, x1) x (xx))) (1o)_1xo-xalee -xllise with - 0 generated a marginal flow th that has the form:
1/1 (xo) Xo + Ce.nlao(x (x) - @ ( (x))
where Ce is a constant that depends on and the dimensionality of the space n.
therefore, for all the points X6 sampled from go that have low likelihood under the final distribution F (so (Xo) << qo (),)), the transform u (Xo) Rs Xo + Ce (Xo) Vx. 91 (xo) will push the sample toward the regions of higher likelihood of q1- Notably, the learned flows also have the property 2/1 (Xo) = Xy when 91.
4 FLEXIBLE DOCKING
In the flexible docking task, our goal is to learn the joint distribution over the bound structures (equivalently, poses thereof) of a protein-ligand ligand complex given the distribution over the unbound structures of the protein. We focus on the pocket-based flexible docking task, but all components of framework be applied to the blind docking setting as well (unknown protein pocket).
In the previous section, we motivated the necessity of Unbalanced FM for the flexible docking task, and the attractive properties that emerge from composing unbalanced flows. In this section, leverage these insights to decompose the flexible docking task into two subtasks, both modeled as unbalanced flows: i) docking over the manifold degrees of freedom (Section 4.1 and ii) structure relaxation (Section
4.1 DOCKING OVER MANIFOLD DEGREES OF FREEDOM
Ligand and protein poses can regarded as elements of 3M and 33,,, where ni and np are the number of atoms in the ligand and protein respectively. However, during docking, ligand flexibility is largely concentrated in the torsion angles at rotatable bonds (Jing et al. 2022)) while for proteins, the flexibility lies in the backbone frames and sidechain torsion angles (Jumper al. 2021). Motivated by the success of Intrinsic Diffusion Models (Corso) 2023) in similar domains, we reduce the space of ligand and protein poses by defining our generative model over these degrees of freedom.

---

Published as a conference paper at ICLR 2025
Transport of distributions formulation. For the distribution over ligand poses, we largely follow DIFFDOCK (Corso et al. 2022). learning a diffusion model over the product space of rotations, trans- lations, and torsions, F SO13) x RR x T To model conformational changes in protein structures upon docking, we employ our Unbalanced FM framework. The prior ạo is defined as the distribution of (computationally generated) unbound structures, while the target distribution ş is defined over the crystallized bound structures. For the prior, given the current inability of computational models to sample the diverse conformational space (Jing et al. 2024), we use outputs from ESMFold (Lin et al. 2022) with the addition of small Gaussian noise. Designing an Unbalanced FM objective then amounts to choosing a coupling (Xx, X1), a conditional probability path pe (x(Xo. X1), (Xu, X1) q-
and the associated conditional vector field u, (x(xxxo, x1).
Choice of coupling q. A key requirement for q is to be able sample pairs during training. Because we typically only have access to one sample for the distribution over bound structures (the crystal structure in PDB, typically unique), we cannot define q via Unbalanced OT. Therefore, we approximate the optimal coupling with the distribution ((xo. x1) о go (Xn)gi (x1 To(x. ) where c(xo, X1) is defined as the aligned RMSD between the Ca positions of the residues in the pocket and its neighborhood, and Clock is an empirically chosen cutoff to balance sample efficiency and mapping complexity. We can sample from q by taking independent samples from 4o and and rejecting if c(xo, x1)
Clock-
Flow Matching on SE(3) and T. For a protein with n residues and my sidechain torsions, we define the flow over the product space SE(3)" SO(2)T², where the SE(3) frame for each residue corresponds to roto-translation around the Co atom, and the hypertorus Tm over sidechain torsions.
Following the disintegration of measures (Pollard 2002), every SE(3)-invariant measure can be broken down into a (0n3)-invrriann -invariant measure and a measure proportional to the Lebesgue measure on IR3 allowing us to build flows independently on SO(3) and R Following Chen & Lipman (2024), given two points (Xo, x1) q, the conditional probability path between Xo and is given by the geodesic between them, x expx. (t logxo (x1)), and the corresponding conditional flow is (x, x1)
log,,
1-t
For (3)) the geodesics can be computed efficiently by using the axis-angle representation (equivalent log(x,)) and the parallel transport operation (left multiplication with xo). while exp is simply the matrix exponential. We view the torus T as the quotient space 1R/2xZ, thus expx, (x1) (xo + x1) mod 2n (equivalent wrapping around RR, and logx, X1 arctan (xsin(x) Xo). cos Xo)).
Confidence discriminator. As discussed in Section 3.1 it is necessary to apply rejection or filtering step at the end of the unbalanced flow. This is because, by the design of the Unbalanced FM framework, some of the samples produced will be incorrect (i.e. far from high-likelihood regions of q1). To do this we train a confidence model (Corso et al. 2022) to predict the likelihood that sample from qi is within 2 ligand-RMSD and 1A AA-RMSD of the input pose. This confidence score is then used rank or select different samples obtained from the flow.
4.2 STRUCTURE RELAXATION
we frame structure relaxation as an (unbalanced) flow, where the prior go is now the distribution over poses generated by the manifold docking process in Section 4.1 and the target distribution 9ı is still defined over the crystallized bound structures.
Motivated by the results from Section 3.3
Choice of coupling 9. We adopt similar definition for as Section 4.1 defined as the average of the aligned atomic RMSDs between the predicted pose for both the ligand and protein (not just the latter) and their ground truth counterparts. We also use smaller cutoff Cdock to enforce small structural changes.
except now c(Xu, x1) is
Crelax <
Flow Matching for Relaxation. We define the flow over the Euclidean space. Given (Xo, x1) q. the conditional probability path (x Xg. is defined as X (1 t) + X1 tt, and the associated conditional vector field nn (x(xp, x1) = x1-x
7

---

Published as a conference paper at ICLR 2025
Energy-based loss. Molecular structures characterized by degrees of freedom with vastly different stiffness, for example the Generalized Amber Force Field (GAFF) is specified such that a single bond between Carbon and Nitrogen atoms allows oscillations with vibrational frequency of 4401 cm- (Wang et al. 2004). This causes even small W2 approximation errors to be particularly problematic along such degrees of freedom. Therefore, to further improve the flow's ability to accurately relax degrees of freedom with narrow distributions, we introduce an additional energy- based loss encourage the model sample low-energy poses.
A similar approach has been explored in the scope of works on Boltzmann Generators (Noé et al. 2019) where models are trained directly to minimize the reverse KL divergence between the generated samples and some energy function by evaluating the energy of the generated poses and their log- likelihood under the model. Our approach distinguishes from this two main aspects: firstly, to preserve the stable and efficient training regime of (unbalanced) flow matching we do not use full sampling trajectories but only evaluate the energy of the xı = x, + (1 - (x;; () predicted poses for & close to I; secondly, instead of using force-field based energy functions that tend to be very unstable on non-equilibrimm poses we employ a handcrafted flat-bottom potential with the following form:
Lenergy Σ max -
x(3) U.i 0) + max (Lis-11s
x 10)
(10)
where Lig and Us, are lower and upper bounds (if any) on the distance between each pair of atoms (i,j) and depend on whether atoms i and are covalently bonded, substituents of the same central atom, or constitute an intermolecular cross term between the protein and ligand. Further details on the formulation are provided in Appendix [B.1
Energy filtering. During inference, we utilize this flat- Bat-bootoom potential combined with the initial poses confidence scores the Unbalanced FM discriminator to select the predictions for which relaxation was successful. Starting with conformers predicted by the flexible docking model, we use the docking confidence model score and prioritize the predictions. After relaxation, we use the flat bottomed potential to select only the relaxed conformers for which the potential is zero. We then choose the relaxed pose with the greatest confidence score. If all relaxed conformers have non-zero potential value, we instead return the highest confidence conformer prior to relaxation.
EXPERIMENTS
Benchmark. We train and test our models on the widely adopted PDBBind benchmark (Liu et al.. 2017p. We use computationally generated structures from ESMFOLD (Lin et al. 2022) as samples from the distribution of unbound structures. We also evaluate on PoseBusters, a recent benchmark dataset curated from the PDB, with several filtering steps and sequence-based clustering. Since we focus on the flexible docking task, we evaluate the accuracy of both the predicted ligand and pocket atom poses. We first align the predicted pocket atoms to their ground truth counterparts, and then compute the heavy-atom RMSD between the predicted and ground truth poses for the ligand (with permutation symmetry correction) and pocket atoms. For the ligand, following Huang al. (2024). report the median and the percentage of predictions with RMSD below 2Å. For the pocket atoms, we similarly report the percentage of predictions with RMSD below 1Å. Finally, report the pose quality using the percentage of top-1 predictions passing the series of tests from [Buttenschonn et al. (2024) referred to PoseBusters checks.
Baselines. For the PDBBind benchmark, we compare FLEXDOCK to well established search-based methods SMINA and GNINA, and flexible ML-based pocket level docking methods in DIFFDOcK- POCKET (Plainer et al.
2023) and RE-DOCK (Huang et al., 2024). All methods receive as input - the ligand with an initial seed conformation (using RDKit). ni) the ESMFold predicted structure, and iii) the pocket residues. Additional details regarding the experimental setup, data, and baselines can be found in Appendix [Dlalong with link to the code used for our experiments. For the PoseBusters benchmark we also report other baselines for pocket- pocket-based docking -tese include both rigid methods
taking in the holo-
like DEEPDOCK (Liao et al. 2019). UNIMOL (Zhou et al. 2023), GOLD (Verdonk et al., (003) and VINA (Trott & Olson 2010) as well as ML-based cotolding methods
UMOL (Bryant et al. 2024) and ALPHAFOLD3 ( (Abramson et al. 2024)
8

---

Published as a conference paper at ICLR 2025
Table 1: Top-1 PDBBind ESMFold Docking Performance We standardize all the generative modeling- based methods to take 10 samples. indicates results missing due to baselines taken from previous work where output structures were not provided. "Note that REDoCK uses slightly different definition of pocket making the results not fully comparable, while their code is not available.
Ligand RMSD
All-Atom RMSD
% PB
% PB valid and
Method
Runtime
% 2A
Median A
% < I
valid
L-RMSD 2
SMINA (rigid)
6.6
7.7
N.A.
258
SMINA
3.6
7.3
5.2
1914
GNINA (rigid)
6.7
7.1
N.A.
93.3
6.7
260
GNINA
8.4
7.9
4.5
91.3
7.7
1575
Di#=Dock-POCKET (rigid)
37.5
3.0
N.A.
25.9
4.9
17
DIFFDOCK-PoCEE POCKET
41.8
2.5
32.4
30.1
5.9
17
REDOCK"
39.0
2.5
FLEXDOCK
39.7
2.5
41.7
72.9
33.7
11
Pasebusters V2
Pocket Residues Specilied
UFN (Cutoff caoa?
INGD 2A
UPM (2A)
P5-Vuld IMSD
UFM (10AA
Orade Rerfommarr
FM
0-
RMBD
Deepbock
UriMo
Unicl
RiolDosk
0OLD
Vha
AF3
Metric
Figure 4: Performance of different pocket-based methods on the PoseBusters v2 benchmark. GOLD, VINA, DeepDock and UniMol take as in- put the holo-structure the protein, while UMol, Flex and AlphaFold3 only the sequence. AlphaFold3 was trained on dataset larger.
Figure 5: Docking performance for different cutoffs Cdock FM corresponds to UFM with We generate 10 samples using small model (4M parameters) without any addi- tional filtering relaxation, and evaluate the best prediction these samples for every complex.
Cdock
Results. Table presents the comparison of previous methods in the field with the overall FLEX- DOCK model which shows improvements in many metrics. To better understand the value of our methodological contributions, it is useful directly compare with DIFFDOCK DHCDO-OTOCCTT as this uses the same architecture and training regime. Here we see large improvements in the pocket all-atom accuracy with AA-RMSD < Á improving from 32% to 42% and the pose quality (measured by the proportion of PB valid poses) improving from 30% 73%.
On PoseBusters, our method performs significantly better than the fast ML-based rigid docking models DEEPDOCK and UNIMOL even without receiving the holo-structure. When compared to large-scale co-folding models, has similar RMSD performance but significantly better validity than UMOL while being significantly faster (average 11s per complex for FLEXDOCK and 206s for UMOl excluding its MSA computation step). ALPHAFOLLD performs significantly better, however, it has been trained on more than one order of magnitude extra data compared to FLEXDOCK and UMOL.
Ablations. Ablations regarding the choice cutoff Clock (See Para. can be found in Figure 5
The comparison with standard FM highlights the importance of unbalanced flows in tackling this task. Using cutoff of 4Å gives better performance across both ligand and protein pose prediction metrics.
We also perform ablation studies in Table 2 to decompose the improvements from the different components of FLEXDOCK. This shows that, the one hand, the improvement in All-Atom RMSI)

---

Published as a conference paper at ICLR 2025
Table 2: Ablation of the effect of the various components of the relaxation unbalanced flow described in Section 4.2 We now split the runtime into the docking time + relaxation time.
Method
% 2
Ligand RMSD
Median À
All-Atom RMSD
% I
PB
L-RMSD < À
% PB valid and
Runtime (s)
FLEXDOCK relaxation)
2.6
FLEXDOCK (OpenMM Relaxation)
38.2
2.6
42.1
68.4
30.3
10 155
FLCXDoCK trelaxation w/out energy loss)
40.9
2.5
44.1
17.3
12.1
10
FLEXDDOK energy filtering)
2.5
0.3
FLEXDOCK
39.7
2.5
41.7
72.9
33.7
10 0.3
can be attributed to the improved modeling of protein flexibility through the Unbalanced FM over the manifold degrees of freedom, not the relaxation. On the other hand, the improvement in pose quality can be largely attributed to the different steps of the relaxation Unbalanced FM which improves the %-PB valid from 13 to 73. Among these, we highlight the critical role of the energy loss. We additionally compare our relaxation to the OpenMM relaxation offered in the PoseBusters benchmark. Our relaxation routine achieves an improvement 4.5% PB-Valid metrics, while being at least order of magnitude faster.
CONCLUSION
While generative models revolutionizing the field of structural biology, existing frameworks lack the flexibility needed to tackle outstanding domain-specific issues. In this work, we presented Unbalanced Flow Matching. generalization of the Flow Matching framework that relaxes the marginal preservation assumption. Theoretically, we showed this approach allows us to improve the approximation error FM incurs at the cost of some sample efficiency. We apply Unbalanced FM to enhance the ability of flexible docking models to model the flexibility of protein structures and predict poses with good energetic properties. Empirically, FLEXDOCK improves these properties over the best existing methods on the PDBBind benchmark.
show that the resulting framework
ACKNOWLEDGEMENTS
We thank Ya-Ping Hsieh, Itamar Chinn, Hannes Stark, and Bowen Jing for their valuable feedback and insightful discussions.
This work was supported by NCCR Catalysis (grant numbers 180544 and 225147), a National Centre of Competence in Research funded by the Swiss National Science Foundation, the NSF Expeditions grant (award 1918839; Collaborative Research: Understanding the World Through Code), Abdul Latif Jameel Clinic for Machine Learning in Health, the DTRA Discovery of Medical Countermeasures Against New and Emerging (DOMANE) Threats program, the MATCHMAKERS project supported by the Cancer Grand Challenges partnership financed by CRUK (CGCATF- 2023/100001) and the National Cancer Institute (OT2CA297463) and the Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger. Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.008797 2023.
Amr Alhossary. Stephanus Daniel Handoko, Yuguang Mu, and Chee- Shhee-Keng Kwoh. Fast, accurate, and reliable molecular docking with quickvina 2. Bioinformatics, 31(13):2214-221 2216,, 2015.
Joe Benton, George Deligiannidis, and Arnaud Doucet. Error bounds for flow matching methods. arXiv preprini arXiv:2305. 16860, 2023.
Patrick Bryant, Atharva Kelkar, Andrea Guljas, Cecilia Clementi, and Frank Noé. Structure prediction of pootein-ligante complexes from sequence information with umol. Nature Communications, 15 (1):4536, 2024.
Martin Buttenschoen, Garrett M Morris, and Charlotte M Deane. Posebusters: Ai-based docking methods fail to generate physically valid poses or generalise to novel sequences. Chemical Science, 2024.
Ricky TQ Chen and Yaron Lipman. Flow matching on general geometries. In The Twelfth Interna- tional Conference on Learning Representations, 2024.
Gabriele Corso. Modeling molecular structures with intrinsic diffusion models. Massachusetts Institute of Technology, 2023.
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock:
Diffusion steps, twists, and turns for molecular docking, arXiv preprint arXiv:2210.01776 2022.
Peter Eastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle A Beauchamp, Lee-Ping Wang, Andrew C Simmonett, Matthew Harrigan, Chaya D Stern, et al. Openmm Rapid development of high performance algorithms for molecular dynamics. PLoS computational biology, 13(7)ze 1005659, 2017.
Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian J Theis. Unbalancedness in neural monge maps improves unpaired domain translation. In The Twelfth International Conference on Learning Representations, 2024. URL https
//openreview net/
Nafisa M Hassan, Amr A Alhossary, Yuguang Mu, and Keong Kwoh. Protein-ligand blind docking using quickvina-w with inter-process spatio-temporal integration. Scientific reports, 7(1); 15451, 2017.
Yufei Huang, Odin Zhang, Lirong Wu, Cheng Tan, Haitao Lin, Zhangyang Gao, Siyuan Li, Stan Li, al. Re-dock: Towards flexible and realistic molecular docking with diffusion bridge. arXiv preprint arXiv:2402. 1459, 2024.
Bowen Jing, Gabriele Corso, Jeffrey Chang. Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35:24240-24253 2022.
Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. arXiv preprint (4XXv:240220044455 2024.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-588)
2021.
David Ryan Koes, Matthew P Baumgartner, and Carlos J Camacho. Lessons leamedd in empirical scoring with smina from the csar 2011 benchmarking exercise. Journal of chemical information and modeling, (193)):893-1999 2013.
11

---

Published as a conference paper at ICLR 2025
Zhirui Liao, Ronghui You, Xiaodi Huang, Xiaojun Yao, Tao Huang, and Shanfeng Zhu. Deepdock: enhancing ligand-protein interaction prediction by a combination of ligand and structure infor- mation. In 2019 IEEE International Conference Bioinformatics and Biomedicine (BIBM), pp. 311-317. IEHE, 2019.
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Allan dos Santos Costa, Maryam Fazel-Zarande Tom Sercu, Sal Candido, et al. Language models of protein sequences at the scale of evolution enable accurate structure prediction. BioRxiv, 2022:500902, 2022.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamul Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXi::2210.02777, 2022.
Zhihai Liu, Minyi Su, Li Han, Jie Liu, Qifan Yang. Yan Li, and Renxiao Wang. Forging the basis for developing protein-ligand interaction scoring functions. Accounts of chemical research, 50(2): 302-309, 2017.
Wei Lu, Jixian Zhang, Weifeng Huang, Ziqiao Zhang, Xiangyu Jia, Zhenyu Wang, Leilei Shi, Chengtao Li, Peter G Wolynes, and Shuangjia Zheng. Dynamicbind: Predicting ligand- protein- complex structure with deep equivariant generative model. Nature Communica- fions, 15(1): 1071 2024.
Felipe Maia Polo and Renato Vicente. Effective sample size, dimensionality. and generalization in covariate shift adaptation. Neural Computing and Applications, 18185):18187-1811 2023.
Andrew T McNutt, Paul Francoeur, Rishal Aggarwal, Tomohide Masuda, Rocco Meli, Matthew Ragoza, Jocelyn Sunseri, and David Ryan Koes. Gnina 1.0: molecular docking with deep learning. Journal of cheminformatics, 13(1):43, 2021.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 147, 2019.
Michael Plainer, Marcella Toth, Simon Dobers, Hannes Stark, Gabriele Corso, Céline Marquet, and Regina Barzilay. Difkdock-pocket: Diffusion for pocket-level docking with sidechain flexibility. 2023.
David Pollard, user's guide to measure theoretic probability. Number 8. Cambridge University Press, 2002.
Aram-Alexande Pooladian, Heli Ben-Hamu, Carles Domingo-Enrinn Enrich, Brandon Amos, Yaron Lipman, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint urXiv:2304. 14772, 2023.
Zhuoran Qiao, Weili Nie, Arash Vahdat, Thomas F Miller III, and Animashree Anandkumar. State- specific prgtein-ligand complex structure prediction with a multiscale deep generative model. Nature Machine Intelligence, pp. 1-14, 2024.
Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne, Aligned diffusion schrödinger bridges. In Uncertainty Artificial Intelli- gence, pp. 1985-1995. PMLR, 2023.
Hannes Stärk, Octavian Ganea, Lagnajit Pattanaik, Regina Barzilay, and Tommi Jaakkola. Equibind: Geometric deep learning for drug binding structure prediction. In International conference machine learning, pp. 20503-20521. PMLR, 2022.
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid fector-Brooke Kilian Fatras, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint arXiv:2302. 00482, 2023.
Oleg Trott and Arthur J Olson. Autodock vina: improving the speed and accuracy of docking with a new scoring function, efficient optimization, and multithreading. Journal of computational chemistry, 31(2):455-661 2010.
12

---

Published as a conference paper at ICLR 2025
Marcel L Verdonk, Jason C Cole, Michael J Hartshorn, Christopher W Murray, and Richard D Taylor. Improved protein- ligand docking using gold. Proteins: Structure, Function, and Bioinformatics, 22(4):609-623 2003.
Junmei Wang, Romain M Wolf, James W Caldwell, Peter A Kollman, and David Case. Devel- opment and testing of a general amber force field. Journal of compulational chemistry, 25(9): 1157-1174, 2004.
Gengmo Zhou, Zhifeng Gao, Qiankun Ding, Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang. and Guolin Ke. Uni-mol: A universal 3d molecular representation learning framework. 2023,
13

---

Published as a conference paper at ICLR 2025
A PROOFS
Note: in all derivations and definitions in this section, we will assume that the distributions we work with are defined in Euclidean space and have full support.
A.1 UNBALANCED FM SAMPLE EFFICIENCY
Proposition 1. The effective sample size, ESS", for sampling qı when having access to samples of go and a perfectly trained flow with coupling distribution is bounded by:
ESS (q) 2 exp - D2(20 (0o)))ls))) - D2(1s, [g:)]
(11)
where D2 is the Rényi Divergence of order 2.
Proof. When comparing pairs of distributions p and q the (population) effective sample size ESS (p. q) is defined as (Maia Polo & Vicente 2023):
ESS (p,9) := exp(-
and it can be considered as the percentage of effective samples from q one can obtain when taking samples from p.
Similarly, we define ESS* (q) in our setting as the percentage of effective samples from qı one can obtain from du using 1/1 While t could be directly applied any distribution, including qo is hard to model its pushforward analytically. On the other hand, we know that (x, is the pushforward of 9x. therefore we can obtain samples from qi by (1) reweighting samples of go into samples of Ixo. (2) transporting samples from (xo, to samples of (x, and (3) reweighting samples from (x, into samples from Q) By assumption of perfect flow step (2) has perfect efficiency, however, steps (1)
and (3) both may require more than one sample in expectation be unbiased. This translates into an efficiency equal to the product of the two effective sample sizes:
ESS*(g) (q) 2 ESS* (qu- (xa) ESS* (gxi 41) = exp 1 D2(gullluull)) (x.))l exp 1-pl-D2(e) D2(xx. 1q1)
cxp - D2(qul D2(4xı
where the inequality derives from the possibility of the existence of more effective procedures for this sampling that do not require passing from samples gxo and 4x1
A.2 COMPOSING UNBALANCED FL.WSS
Proposition 2. Assuming (o and qn are densities with a bounded Lipschitz constant. A coupling q(xo, x1) xx 9o (xo)q1 (x1 S xillse with e - generates a marginal flow V/2 that has the form:
t (xx) == No + Ce. (x)) V1x (xxo) - 91 (xx) (ooo))
(12)
where Cep is constant that depends on and the dimensionality of the space n..
Proof. Let (xx) be the vector field that generates the flow: (x) = u (x)). By definition
ue ue(x;) is the marginal vector field that optimizes the objective function of the defined Unbalanced FM. Therefore its derivative must be zero:
V (x (xxxxx))~ (xx) us (x(xx(/ss )|1
V (x E xxx (xx) liz (xx) (x1 x
Let 8 = X1 Xo. Expressing the expectation as n-dimensional ball integral:
J101ss 290 (x: 8x)q1(x( + 5(1 ()x - 66) Zg d8
14

---

Published as a conference paper at ICLR 2025
Therefore:
u, (x() = { /1s
(x0 (x, + 0( - t))
d8
To approximate this integral as - expand go and qı using Taylor series expansion around X₃.
qo(x: öt) gx((x) t8 (x)
qi (x, 0(1 - t)) qı (xx) (1 t) Vig (xx)
Substituting the approximations into the integral:
(x() 2 f S8
[((0(xx)) t6 (x(xx)) (1 t)ö (xi(x)))
Expanding the product inside the integral:
(x())) Rs / Jil (x)))))11)))) + (x(x)) - l)ö (x. 15 (x)())
Higher-order terms involving (6 Vqu) Vqn) will be of order 2² and can be neglected as c Since the integral is over symmetric ball, terms with odd powers of vanish. Hence, only the cross terms contribute:
u (x1)
8 [[[1 (x,6 (x,) ton (x (x0(
S181is
Using the symmetry and volume element:
J 1
T
8,6, dd
8iz
[mm/2 ( n n
Therefore:
nn
n/2
(x(x
n + 2 2n2m2 1) [(1 (xzx tęg (x;,
(oo
To determine the residual transformation brought by the full marginal flow, we need to integrate the vector field (x,)) over the path. Integrating over the path challenging, but since have restricted the flows to be smaller than e, the Lipschitz constant of u, is bounded (because the Lipschitz The change in u, within the flow will
constants go and qi are bounded), and we taking be negligible. Therefore, we can approximate the marginal flow as:
n/22
vi (xx) ~ Xo +
1= () [[1 t)qu(xo) Vqr Vg (xo) - tqi (x)) V (xo)) dt
Evaluating the integral:
n22
2
(1-
1n (xx)) so No + (n + ( (n/2
(x))
(x V (xx)
2
n n/2 m22 n+2
1i (x(x)) x Xy + 2(n + 2) (n/2 1)
(20(x))VQ1 - (xo)
(go(xo)) (xx))xx((x)))
15

---

now proceeds Published as a conference paper at ICLR 2025
B TRAINING AND INFERENCE
In this section, present the training and inference procedures for our manifold docking (Algorithms 3 4)n and relaxation models (Algorithms ()d (nd(l) We refer to unbound protein structures as apo structures, and the bound structures holo structures. Recall that our goal is to learn a distribution over holo structures, given the apo structure and seed conformation of the ligand. Similar to (Corso et al. 2022), are in a setting, where traditional generative modeling where one has access to multiple samples from the same data distribution, we only have a single (x",, yapo Ynooo per protein-ligand complex. This implies that the training loop (Algorithms different distributions, along with a single sample from that distribution. This sample is then accepted or rejected depending on the cutoff Clock, thus inducing unbalanced coupling and flow.
Algorithm 3: TRAINING EPOCH: MANIFOLD DOCKING
Input: Training Pairs {(x", Y npos Yholo J: RDKit predictions (c), RMSD cutoff Cdock Input: Pocket radius r; Pocket Buffer b
Input: Ca operator (:lca
foreach x"x apo: Vhoio do
Let Xo f- arg minly RMSD(x" x)
Voenter: {i] pocket EXTRACTPOCKET(CS apo Ybolo T, b) Y Yapo +- R.MSDALIGN(Y appo: Yholo pocket)
if R.MDD(Y apo: Yhoio) Cdock then
continue
else
apo: Rbb + FRAMEALIGN(Y opo Sholo)
FA
y
y Yapo
SC
A +- SCCONFMATCH( apo Ynol)
Sample ~ 1(
// Ligand Diffusion
Sample Ar, AR, AO from diffusion kernels xf* (-10)
Compute x, by applying (Ar, AR, A0) to X
// Protein Flow
bb 150 -or bbb COMPUTETIME(T, a a Op
Interpolate Art [yapo] Cca (1 - + [y hooo
Ca
jr.b (-12) + [yhoto] - [[ypoo] CC Ca
Interpolate A n2+b + exp (1ab log( Ppbb))
1op R (ART
pubbb (-1z) +
Ut
1-
Interpolate AOt & exp (t ui (-1(=) loRaep (AI")
I-t4
Compute y: by applying (b A R6 A to Yapo
Predict scores and drifts a, 3,,,,,,,,, 1 - Stx1)
// liiandd Loss
Lig ||a: V log (o (- + ||3 Vlog pf 1o0ppt((((1100 (-10)112 || V (. (. (0)
// Protein Lcss Epor ||8 u. rnhh (11=) + - w. (-12)112+ ||m 100-2((1122 (+121111
Apply optimization step on Lprot + Lug
end
end
Pocket Extraction. As our focus is on the flexible protein docking task, we first extract the protein pocket given apo and holo-structures. We define the pocket residues all residues in the holo-
16

---

Published as a conference paper at ICLR 2025
Algorithm 4: INFERENCE: MANIFOLD DOCKING
Input: RDKit predictions {c}, Apo structure Yapo of the protein pocket
Input: Inference Steps N
Sample N ~ ~ Ry ((S((3))) An ~ o()
Apply ON Rn to get XN
Set YN < Yapo
AL 1/N
for 7 + to do
1 <-n/N
Predict scores and drifts a, a,8,7 8,7,0, E, t S(xn.
// Ligand Updates
Ang (N/NN) - ( 1)/N)
Ao mm (n/N) o
1)/N)
Ao Te (n/N) - (( 1-nn-11 1) [/N
Sample Zuy 2oo Zur from (O, (0, (,, o
Apply (a,8,7) 8,7) to X, to get xn-1
11 Protein Updates
Anth + 8- At
ARth + E. At
Ags {- At
Apply (Arhe AR AP.)) Yn to get Vn-1
end
structure that have least one heavy atom within 5Á of any ligand atom. Given these pocket residues, the pocket center is defined based on the positions of the Ca atom in the apo structure. To construct the geometric graphs (Appendix D). we also use the residues which have a Co atom within 20Å of the pocket center. This additional buffer is added improve the model's 's robustness exact pocket definitions, and also add geometric information from the pocket neighborhood.
Aligning Apo-Holo Frames. A residue frame (Jumper et al. 2021). is characterized by a tuple (R, I) F3),, where the rotation R is about the origin of the residue, and specifies the position of the Ca atom. Before applying the conformer matching step (explained below) to the protein sidechains, we align the frames of the apo and holo structures, by computing the rotation that aligns the N - Ca vectors of the corresponding residues. The alignment will not be perfect owing to differences in the bond lengths and bond angles between the computationally generated and ground truth structures, but provides the closest modification of the apo structure backbone to the holo-structure one.
Conformer Matching. For both the ligand and the protein sidechains, we apply the conformer matching procedures in (Jing et al., 2022) and (Plainer et al. 2023), where, given the local structures from computational methods, we find the closest (in an RMSD sense) structure to the ground truth by modifying the appropriate torsion angles. The conformer matching procedure is employed to prevent distribution shift between training and inference in the local structures that are considered rigid in the manifold docking process. To elaborate. the local structures (such as bond lengths and bond angles) vary between RDKit (for ligands) and FSMFold (for proteins) generated structures, and their ground truth counterparts. If we train models with ground truth local structures, this would cause a distribution shift at inference time, when we only have access to local structures, provided by RDKit and ESMFold.
B. ENERGY LOSS FORMULATION
To construct the flat-bottom potential defined in Equation 10 for any pair of atoms (i, j), we define the lower and upper distances bounds Lij and U., in the following manner:
1. If the atom pair (i. represents covalent bond between ligand atoms
17

---

Published as a conference paper at ICLR 2025
Algorithm 5: TRAINING EPOCH: RELAXATION
Input: Training Pairs {Kooock Xbolo ddock Ybolos pooket)) RMSD cutoff Crelax
foreach
Xdock, Xholo: Moolo:-dockk Ydocks Yhoto: {e) pocker do
if 0.5 RMSD(y Yholo) 0.5 RMSD (Xaooks Xholo) Crelax then
continue
else
Sample 1((01)
// Ligand Flow
Interpolate Art lig + Xdock (1 1) + Xholo
Ut
lig
(-12) +- Xholo Xdock
Apply Art to ddck get X
// Protein Flow
Interpolate Arp f- Ydock (1 (1-t) + Yholo
u. euo (-12) Yhalo Vdock
Apply Art prot to yaack to get Yt
// Make Prediction
Predict drifts a f- 8, (xt, Vt, Apply a to Xdock to get X
Apply 3 to Vdack to get y1
81 f RMSDALIGN &1, V1, Ynolor pockes)
1/ Compute Loss N Dig N Ei |x
(1)
(1)
12
holo
Cpor NI
E, 11y
(8)
Ynolo
(i)
1|2
Linternal_energy energy
Eis max
{ (i)
x (1) U2j 0) max (Ie. - (x(0) x .0)
Leross, Leress_energy energy
Eis max
[Tij
x
-
L + ulom Corom + X₁
energy
(Cimornal + Ccroas_cnorgyy)
end end
Lig is set to 0.75 times the lower distance bound used by RDKit
U is set to 1.25 times the upper distance bound used by RDKit
2. If the atom pair (i. j) represents substituents of the central ligand atom
Tr.j set 0.75 times the lower distance bound used by RDKit
UtJ set to I 1.25 times the upper distance bound used by RDKit
3. If the atom pair (i. j) represents two ligand atoms that are neither covalently bonded nor bonded to the same central atom
Li.j set to 0.8 times the lower distance bound used by RDKit No upper bound is used
4. If the atom pair (i. j) includes one protein and one ligand atom
Li,j set to the sum of each atom's 's Van der Waal's radius used by RDKit minus 0.75
No upper bound is used
5. If the atom pair (i.j) represents two protein atoms
No lower bound is used
No upper bound is used
Under this formulation, any conformation for which the potential is zero will pass the following PoseBusters validity checks defined in (uttenschoen al. (202)):
18

---

Published as a conference paper at ICLR 2025
Algorithm 6: INFERENCE: RELAXATION
Input:
Manifold Docking Predictions (Xdock. Ydock)
Input: Inference Steps IV
Set X1 < Xdock
Set Y1 Ydock
At < 1/N
for 7 OtoN do
1 < n/N
Predict scores a, + s(Xn;
Apply a to Xn to get Xn:
Apply 8 to Yn to get Yni
end
Bond length validity due to condition (1)
Bond angle validity due to condition (2)
Internal steric clash due to condition (3)
Minimum protein-ligamd distance due to condition (4)
During training, we scale the energy loss term from weight / = 0 to I at / with the exponential schedule: x e and a = 225. Additionally, since the RMSD of the ligand typically greater than that of the protein in docked poses, we further scale the ligand and protein losses with weights Aog 1.0 and Arro respectively.
C MODEL ARCHITECTURE
We use message passing networks based tensor products of irreducible representations (irreps) of SO(3), implemented with the e3nn library.
Graph Construction. We represent structures geometric heterogencous graphs, with nodes comprising ligand heavy atoms, receptor residues in the pocket and neighborhood (located at the position of Ca atoms), and the heavy atoms of the pocket residues. We chose to only model the heavy atoms of the pocket residues for two reasons i) this provides a useful sparsity constraint for computational and memory efficiency, and ii) typically, most of the conformational changes in the protein involve the pocket atoms, and modeling this explicitly would facilitate downstream applications such affinity prediction. We also adopt different cutoffs depending the types of nodes being connected, largely following (Corso et al. 2022):
Ligand atoms logan- atoms, receptor atoms-receptor atoms, and ligand atom-receptor atom interactions use cutoff of 5Å. Covalent bonds between ligand atoms are explicitly modeled with initial edge embeddings to reflect the type of bond. For receptor atoms, we limit the maximum number of neighbors to 12.
2. For receptor residue interactions, we use a distance cutoff of 15Å, with a maximum neighbor limit of 24.
3. For interactions between ligand atoms and receptor residues, unlike (Corso et al. 2022). we found using the dynamic cutoff based on the ligand translation noise cause NaNs during training, possibly due missing connections. We thus used distance cutoff of 80À between ligand atoms and receptor residues.
4. Receptor pocket atoms are also connected to their corresponding residues.
Featurization. We adopted the same featurization DIFrDOCK, using the residue type and the embeddings with ESM2 Language model for the residues, the atom type, and other chemical properties for the ligand and receptor atoms. The relaxation model additionally uses the Van der Waals radii as features for ligand and receptor atoms as well as the RDKit upper and lower distance geometry bounds as features for ligand edges.
19

---

Published as a conference paper at ICLR 2025
Manifold Docking and Relaxation We retain the core architecture of DIFFDOCK (Corso et al. 2022). with the tensor product convolution- based message- message-passing layers, followed by a convolution with the center of mass to predict the rotational and translation scores for the ligand. For the torsion angles in the ligand and sidechain torsion angles in the protein, we use the pseudotorque layer from (Ting al. 2022), adapted accordingly for the sidechains. To predict the rotation and translation flows for the residues (which are (3) equivariant), we use a linear layer that transforms the irreps of the single odd and even vector (one for each flow). Since the residues constitute a coarse- -gaiind representation of the protein, we sum the odd and even vector representations to obtain the predictions. For the relaxation model, we predict single vector for each ligand and protein atom. The magnitudes of the predictions are then adjusted with an MLP.
residue embeddings
Confidence Model. The embedding layers for the confidence model follow the same architecture as the manifold docking and relaxation. The aggregated ligand, receptor residue, and receptor atom embeddings are concatenated, and updated with an MLP to predict the final confidence (a SE(3) invariant output).
D EXPERIMENTAL DETAILS
The code used for running our experiments can be found at Flexdock
Data.
For training our models, we use the BDB dataset (Liu et al. (0117) whose complexes were
extracted from the PDB. Following (Stärk et al. 2022 Corso et al. 2022), we adopt the time-based split of PDBBind, where the 17k complexes before 2019 were divided into training and validation sets, while the 363 complexes after 2019 form the test set. We download the PDBBind data as is provided by EquiBind (Särk et al. (2022). These files are first processed by PDRF xer from the OpenMM toolbox (Fastman et al. (017), to replace non standard residues and add missing atoms. We then used the PDBF i ixer processed files to extract the protein sequence, and predict its structure with ESMFold (Lin et al. 2022). The ESMFold generated files are also processed by PDBF i xer to add missing atoms such as terminal oxygens, at the end of chain. These processed files now constitute our apo structures, while the processed analogues from PDBBind constitute our holo structures. We further remove hydrogen atoms while aligning the apo and holo structures.
To measure the validity and physicality of generated poses, we also use the PoseBusters suite of tests (Buttenschoen et al., (024), which evaluate i) chemical validity, ii) intramolecular validity for the ligand poses based on energetics and geometry, and iii) intermolecular interactions and steric clashes between protein and ligand.
Metries. To evaluate the generated ligand and protein pocket poses, we compute the RMSD between the predicted and ground truth poses after alignment. This alignment is computed based the Kabsch alignment between the atoms in the protein pocket, in the ground truth and predicted poses. To account for permutation symmetries in the ligand, we use the symmetry-correcteo
RMSD of sPyRMSD. For the ligand, besides the median RMSD, we report the % of RMSDs below 2Å. which is a commonly adopted metric for judging the quality of docking predictions (Alhossary et al. 2015 Hassan et al., [017 McNutt et al., 2027). For the protein pocket atoms, besides the median RMSD, report the % of RMSDs below IA, where we chose the I cutoff, typically treated as atomic accuracy.
Training Details. For our manifold docking model (75.3 M parameters), we use an exponential moving average of weights (EMA) during training, which is updated every optimization step, with a decay factor of 0.999. We train the model on 4 RTX A6000 GPUs, with a batch size of 4 per GPU. Every 10 epochs, we run inference for 20 steps with the EMA weights on 500 complexes in the validation set, and save the model with the largest percentage of ligand RMSDs 2A. The initial learning rate of the model is 0.001, which is updated with learning rate scheduler with decay 0.7 if the percentage of complexes with ligand RMSDs < 2Å does not improve over 30 epochs. We train our model for 600 epochs, after which we did not observe noticeable increase in ligand RMSDs 2 metric. We use the ADAM optimizer for all models.
20

---

Published as a conference paper at ICLR 2025
Posebusters V2 Checks
AF3
FlexDock
Dif/Dock
80
16 of passing examples
60
40
20
File loads
Sanitization
Molecular Formalla
Tetrahestra al Chirality
Double Bond Stereschemistry
Bonds
Bond Lengtha
Band Anples
Planar Aramatic Rings
Flanar Double Bonds
Interal Steric Clash
Minimum Protein Ligand Distance
Minimum Distance Organic Cofactora
Energy Ratio
Minimumm Distance to Inarganic Colactors
Volume everlap with prateins
Volume overiap with organic cotactors
Volume everiap with inorganic cofactore
Figure Detailed comparison of the validty checks offered in PoseBusters V2 benchmark
For the confidence model, to generate 20 poses (both ligand and protein) per training complex. For the ligand, we assign label the RMSDS between predicted (after alignment) and ground truth pose is <2A, while for protein pocket atoms, we adopt an RMSD cutoff of IÅ. We train the confidence model for around 100 epochs, and save the model with the best accuracy. We found the model predicting only the ligand pose confidence to offer the best tradeoff between ligand and pocket atom prediction confidence.
smaller version of the manifold docking model 4 M parameters
Runtimes. Similar to other MI. docking baselines, measure runtimes for the manifold docking and confidence model. These runtimes are calculated on single RTX A100 80GB GPU, with the preprocessing steps entailing ESM2 embedding generation and RDKit conformer generation. The geometric graphs are generated on the fly as part of the model and thus are already included in the runtimes.
E ADDITIONAL ANALYSIS
COMPARISONS ON POSEBUSTERS BENCHMARK
We include stratification of different validity checks adopted in PoseBusters, comparing against AL.HAAOLLD3 and DIFFDOCK in Figure The numbers for ALPHAFOLLDS and DIFFDOCK were taken from (Abramson al. 2024). We find that FLEXDOCK generates highly valid molecules with very low computational overhead ( 0.3s), and has significant improvement from DIF:DoCK on the intermolecular validity checks.
E.2 COMPARING RMSD DISTRIBUTIONS
Training and Test Distributions Figure we plot the all alom (AA) RMSDs between ESM- FOLD and holo structures in the training and test sets. While the AA-RMSDs is largely within 4A, there is very long tail of significantly larger distances which impact the expected squared coupling distance when using Flow Matching.
Large Conformational Changes. In Figure 8 we plot the AA-RMSDs between the ESMFold and holo structures, and also between the corresponding predicted and holo structures for large conformational changes (> 4Á). We find that FL.EXDOCK is able shift the distributions towards
21

---

Published as a conference paper at ICLR 2025
Training Tast
0.30
apo_hole hep__hln_en__msot aa_mmass
0.5
0.4
Densty
0.20
0.15
0.1
0.05
0.0
15
25
8
10
RMSDs
Figure 7: Distributions of AA-RMSDs between ESMFold and holo structures in the training and test sets.
Figure 8: RMSD Distributions for Large Confor- mational Changes (>4Å)
lower RMSDs compared to the initial distribution. This is reasonable since the model updates the apo structure to the closest structure based on the learned underlying transport. Additionally, we find that 77.5% of poses pass the Posebusters validity checks, indicating the ability of our model to find physically meaningful poses, even this extreme setting.
ESMFlow vs ESMFold In Figure g plot the AA-RMSDs between the apo and holo structures, where the apo structures are generated either using ESMFold or its flow matching counterpart ESMFlow (Jing al. 2024). The distributions look largely similar, with most of the RMSDs within 5A, followed by a long tail of larger RMSDs. This also highlights the limited ability of current methods capture the diversity in the conformational space of apo protein structures.
All atom RMSD
all rmsds
min rmsds
ESMFlow
ESMFold
0.35
0.30
0.30
0.25
0.25
Proportion
Proportion
0.20
0.20
0.10
0.10
0.05
0.05
0.00
0.00
10
20
30
40
10
20
RMSD (A)
RMSD (A)
Figure 9: AA-RMSD distributions between apo and holo structures, with apo structures generated
using ESMFold or ESMFlow
E.3 SAMPLE EFFICIENCY
In practice, rather than the actual sampling efficiency, one is often interested in indirect measurements of the same, and using that compare different methods. These indirect measurements are based
22

---

Published as a conference paper at ICLR 2025
40
40
samples
samples
5 samples
samples
10 samples
10 samples
30
30-
Lig A0--RMSDD 2A
MARMSDS 1A
20
20
10
5
2.0
4.0
6.0
10.0
FM=
2.0
4.0
6.0
10.0
FM=
Cutoff
Cutoff
Figure 10: Sample Efficiency for different cutoffs Gdock
on performance achieves on desired metrics, with fixed number of samples. In Figure TO plot
the best achievable performance for each Cdock for 1, 5, 10 samples. Comparing FM with UFM (4A), where Cdock 4.0 for instance, we can clearly see that even with 10 samples, FM achieves a worse performance on ligand docking, compared to UFM (4A), with 5 samples, making UFM (4A) more sample efficient.
23

---

Published as a conference paper at ICLR 2025
F VISUALIZED EXAMPLES
(a) 6FE5, L-RMSD: 0.66, AA-RMSD: 0.72, APO-HOLO-RMSD::30 3.93
(b) 60D6, L-RMSD: 0.61, AA-RMSD: 0.81, APO-HOLO-RMSD: 1.67
(c) 6BQD, L-RMSD: 0.51, AA-RMSD: 0.42, APO-HOLO-RMSD: 0.56
(d) 6ROT, L-RMSD: 0.71, AA-RMSD: 0.57, APO-HOLO-RMSD: 0.94
(e) 600Y, L-RMSD: 3.97, AA-RMSD: 8.75
(f) 6CKL, L-RMSD: 3.51, AA-RMSD: 9.01
Figure 11: Visualized predictions for the complexes 6FE5, 60D6, 6BQD, 6ROT, 600Y, and 6CKL in the PDBBind test dataset. L-RMSD and AA-RMSD denote Ligand RMSD and All-Atom RMSD respectively, after aligning the predicted and ground truth pockets. Color Code: Holo protein (in lime), predicted protein (magenta), true ligand (cyan) and predicted ligand (orange).
24


Reviews:
[{"review_id": "Ugn9UloBAN", "reviewer": "Reviewer_sBsN", "summary": "The paper presents a new approach named Flexdock for flexible molecular docking and structure relaxation.  The model is based on Unbalanced Flow Matching (UFM), which relaxes the marginal constraints of FM through a more careful chosen coupling distribution and reduces the modeling complexity. This framework aims to address limitations in existing molecular docking techniques, particularly concerning protein flexibility and the generation of nonphysical molecular poses. From the experimental results on PDBBind benchmark, the proposed method well achieved this goal, especially in generating energetically favorable poses.", "questions": "(Have already stated in the weakness section)", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* As far as I know, the introduction of unbalanced FM in flexible molecular docking is novel.\n* The authors analyze sample efficiency and approximation error with mathematical foundations.\n* The introduction of energy-based loss for relaxation is reasonable and seems very useful in generating energetically favorable poses.", "weaknesses": "* The introduction section lacks necessary citations. For example, in line 43, “Deep learning methods … often force the model to relearn protein folding”. What does “relearn” mean?  Please provide citations for such related work.\n* In the related work section, the authors pointed out some related work but didn’t discuss the difference / advantages of their proposed method over related work.\n* In methods section, the introduction of Eq3 and Eq4 is unclear. What’s their connection and how are they applied in practice?\n* What is the influence of the choice of threshold c_{relax}?\n* Have the authors tried using traditional (non learning based) energy relaxation methods like Rosetta as a step of post processing for baseline models and the basic FlexDock without proposed relaxation? I’m curious how it will perform compared to the proposed relaxation method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5fcF9rLrEF", "reviewer": "Reviewer_7QfK", "summary": "This paper extends Unbalanced Flow Matching (UFM), a novel approach to molecular docking that extends Flow Matching (FM) to better handle protein flexibility and generate physically realistic poses. The main innovation is allowing trade-offs between sample efficiency and approximation accuracy in the transport between distributions. The authors implement this in FLEXDOCK, which chains two UFM components: manifold docking for approximate pose prediction and structure relaxation for refinement. On the PDBBind benchmark, FLEXDOCK improves both docking accuracy and the proportion of energetically favorable poses compared to existing methods.", "questions": "1. Is the manifold docking module capable of blind docking, or is it limited to pocket-based docking only? (This is not a request for additional experiments, but rather an inquiry about the capability of this method.)\n2. How sensitive is the method to the choice of hyperparameters?\n3. What insights do the authors have regarding previous ML-based approaches performing poorly on the PB-Valid metric? Based on the experimental results from diffdock-pocket (rigid), rigid docking methods still generate non-physical structures. Which component predominantly drives the PB-Valid metric in your experiments - intramolecular or intermolecular validity?\n4. Compared to PB-Valid metrics, FLEXDOCK demonstrates no significant improvement in RMSD-related metrics. While RMSD remains a crucial indicator for evaluating approximate poses, one must question the necessity of pursuing validity and physicality when RMSD values are still low and approximate pose estimation lacks precision.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper attempts to address two critical limitations found in previous docking methods: the rigid assumption and the generation of non-physical structures.\n2. The extensive experiments demonstrate the effectiveness of the proposed method. This paper demonstrates significant improvement in the PB-Valid metric, particularly with the simple yet effective energy loss.\n3. The paper analyzes and formulates the flexible docking problem from the perspective of distribution transport, and extends the Flow Matching method specifically for this problem, supported by theoretical derivations.", "weaknesses": "The proposed method and experiments still assume prior knowledge of the binding pocket; despite this known pocket assumption, FLEXDOCK shows no significant improvement in RMSD-related metrics. This suggests that the manifold docking module may not make substantial progress in resolving the approximate pose compared to previous works.\nHowever, this limitation appears to be inherent to the challenging nature of the protein-ligand docking problem and reflects the current state of the field, rather than diminishing the paper's existing contributions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qlmrfgVW9S", "reviewer": "Reviewer_R8VH", "summary": "The authors propose FlexDock, a flow matching-based method for flexible small molecule docking. To address the problem of the limited availability of the bound protein structures, authors formulate Uncoupled Flow Matching, in which they tradeoff the approximation quality of the target distribution and the amount of samples needed from the prior distribution. Authors show that FlexDock outperforms other flexible docking methods in various metrics.", "questions": "1. It would be interesting to look at some examples of the model predictions (overlaid with the ground truth).\n2. It would be helpful to provide the distribution of RMSD on $C_{\\alpha}$ positions between apo and holo structures in the training and test data.\n3. What is $q_0$ for ligands and how does the Unbalanced FM framework combine with the diffusion model for ligands?\n4. What is the actual sampling efficiency of the final model? I.e. how many samples are required on average until the one exceeding the confidence threshold is obtained? Does the Runtime in Table 1 account for it? It would be interesting to look at the dependency of ESS on the selected $c_{dock}$ cutoff.\n5. It has been recently shown that the conventional train/test splits of PDBbind (that authors use in this work), have a significant data leakage [1]. Is it possible that majority of 41% successful predictions (in Table 1) overlap with the training data?\n\n[1] Corso G. et al. Deep confident steps to new pockets: Strategies for docking generalization //ArXiv. – 2024.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed Unbalanced FM framework is a very neat idea and can be very useful in many different applications where the data from one of the distributions is limited\n2. FlexDock demonstrates impressive performance across multiple metrics", "weaknesses": "1. In the motivation for the unbalanced flow matching authors say that due to the limited availability of the data points from $q_1$ (i.e. holo structures), the expected length of the flow will be large and the expected OT cost between $q_0$ and $q_1$ will be high. In my opinion, that would be true if there was a certain level of diversity in samples from $q_0$ (i.e. different apo structures) that correspond to a single (available) sample from $q_1$. But is it really the case for protein structures? I have a feeling that the diversity in apo structures of the same protein is not that high. I believe that the proposed method indeed improves learning the transport, as demonstrated in Figure 4. But I think it would be very beneficial to dissect this experiment and look at the flow distances and transport costs for different cutoffs, or even just plot these distributions for your training data. In my opinion, such a deeper analysis will strengthen the motivation a lot.\n\n2. Authors suggest to use the coupling of the form $q(x0,x1) = q_0(x_0)q_1(x_1)\\mathbb{I}_{c(x_0, x_1)<cdock}$. \n* 2.1 Does it mean that some examples where conformation change upon binding is significant are completely discarded from the training?\n* 2.2 If 2.1 is true, how does the model perform on similar examples (i.e. the ones that change a lot upon binding) from the test set?\n* 2.3 Can it be true that \"All-Atom RMSD% < 1A = 41.7%\" means that around 40% of the test set are proteins that almost don't change upon binding, and the model (trained on couplings with rather strict $c_{dock}$) just doesn't change these structures meaning that it \"recovers\" holo conformations for these 40% correctly?\n\n\n**Minor comments:**\n1. Lines 40-45: could you please provide the references to the works you’re discussing here?\n2. Are the protein states over the trajectory physically-realistic?\n3. Could you please provide a more detailed description of y-axis in Figure 4 (include \"% of passed predictions\" or something like this)?\n4. Line 273 - 274 -> I could not find the introduction of the coupling $q(x0, x1) = q(x0)q(x1)\\mathbb{I}$ in the previous section. What are you referring to?\n5. Line 182 – I am not sure if it is a typo or maybe I am wrong, but to me it looks like (3) is the lower bound of (2).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rzkC6eT5gW", "reviewer": "Reviewer_YA9F", "summary": "This papers introduces Unbalanced Flow Matching (UFM), an extension of traditional flow matching, particularly tackling molecular docking problems considering protein flexibility. While traditional molecular docking methods assumed rigid structure for proteins limiting the accuracy, sometimes failing in scenarios including conformational changes of protein. Based on UFM, the authors propose the framework FLEXDOCK that considers flexible docking, and structure relaxation for optimized pose.", "questions": "1. UFM in extreme cases\n\nJust curious, does UFM succeed even in extreme scenarios where the unbound (apo) and bound (holo) protein structures highly differ?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. A new paradigm coined Unbalanced Flow Matching\n\nThe proposed objective is an extension of the traditional flow matching. Relaxing the strict marginal conservation allows more efficient and accurate mapping between distributions.\n\n2. Two subtasks for flexible docking\n\nBy dividing flexible docking into two subtasks, manifold docking, and structure relaxation, the method results in more physically realistic structures.\n\n3. Empirical results\n\nFLEXDOCK outperforms prior methods on the PDBBind benchmark. Noticeably, this has been done using the same training data and model architecture as DiffDock-Pocket. Ablation studies also verify each component of the proposed method.", "weaknesses": "1. Flexibility and complexity tradeoff\n\nWhile UFM has strength in flexibility compared to prior works, it inevitably seems to face more complexity than before. Could the authors briefly summarize the complexity compared to the algorithms in the prior works, if possible?\n\n2. Choice of coupling\n\nUFM seems to rely much on the coupling distribution. The authors have empirically chosen cutoffs, but this dependency may have problems for generalization. I may have misunderstood the part of choosing the coupling.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nBiI2ej5bG", "reviewer": "Reviewer_ZfZb", "summary": "This paper proposes Unbalance Flow Matching and demonstrates its application to flexible docking, and enhances its performance from 30% to 73%. For those who have attempted the traditional flow on flexible docking or protein domain structural changes, they must have encoutered the same problem addresses in this paper. It offers a good solution backed by theoretical insights and supporting proofs.", "questions": "- In Section 3, the authors mention challenges in applying flow matching directly to flexible docking:\n    \n    > In fact, although the different conformational states of the protein do not change significantly upon ligand binding, their relative weights are often notably\n    altered.\n    > \n    \n    What is the **relative weights** meaning?\n    \n- The proposed second stage is to perform the relaxation by deep learning model. Maybe it is to give an end-to-end solution for more precise docking, as shown in Table 2. But I strongly recommend to conduct the minimisation (relaxation) using force fields instead as what is done in PoseBusters (Section 2.5 and 3.3 in PoseBusters paper). The comparison between PB valid and runtime consumption is preferred.\n- For the apo structure, as mentioned in 4.1: \n>given the relative inaccuracy of computational models to sample multiple protein conformations.\n> \nI’m wondering why the ensembles generated by, for example, AlphaFlow, is inaccuracy. Since it is sampled from the MD trajectories and, in principle, it should provide more accurate structures than the addition of small Gaussian noise on apo structure.\n- The unbalanced pattern is pinned across the training. However, is there a potential risk that, for example, there are some apo structures all have RMSD exceeding 4Å from the holo structures? Under such situation, they are all been filtered out and the sampling would fail. Should the traditional flow and unbalanced flow been given a ratio during the training process (like 5/5)?\n- The comparison between AF3 on PB is preferred (extended Data Fig.4).", "limitations": "", "rating": 10, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- It proposes the Unbalance Flow Matching (UFM), which extends the traditional flow matching with the coupling q by performing some filtering (e.g., RMSD) in practice.\n- The proposed two-stage flow matching process, consisting of a rough docking stage called “manifold docking” and a more precise “relaxation” stage, effectively boosts the PB valid benchmark to 73% while maintaining reasonable computational costs.", "weaknesses": "- I do not see major weakness in this paper, but have several questions as follows.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)