Multi-domain Distribution Learning for De Novo Drug Design
Authors: Arne Schneuing
Ilia Igashov
Adrian W. Dobbelstein
Thomas Castiglione
Michael M. Bronstein
Bruno Correia
Venue: ICLR 2025 Poster (2025)
Keywords: Drug Discovery
Flow Matching
Markov Bridge
Equivariance
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: We introduce DrugFlow, a generative model for structure-based drug design that integrates continuous flow matching with discrete Markov bridges, demonstrating state-of-the-art performance in learning chemical, geometric, and physical aspects of three-dimensional protein-ligand data. We endow DrugFlow with an uncertainty estimate that is able to detect out-of-distribution samples. To further enhance the sampling process towards distribution regions with desirable metric values, we propose a joint preference alignment scheme applicable to both flow matching and Markov bridge frameworks. Furthermore, we extend our model to also explore the conformational landscape of the protein by jointly sampling side chain angles and molecules.

Published as a conference paper at ICLR 2025
MULTI-DOMAIN DISTRIBUTION LEARNING FOR DE NOVO DRUG DESIGN
Arne Schneuing" Ilia Igashov Adrian W. Dobbelstein', Thomas Castiglione", Michael Bronstein & Bruno Correia "cool Polytechnique Fédérale de Lausanne, VantAA, Inc., University of Oxford, 4 "Aithyra {arne. schneuing, ilia. igashov, bruno brr.o.correiaae@eff..c.
ABSTRACT
We introduce DRUGFLOW, generative model for structure-based drug design that integrates continuous flow matching with discrete Markov bridges, demon- strating state-of-the-an performance in learning chemical, geometric, and physical aspects of three-dimensional protein-ligand data. We endow DRUGFLOW with an uncertainty estimate that is able detect out-of-distribution samples. To further enhance the sampling process towards distribution regions with desirable metric values, we propose a joint preference alignment scheme applicable to both flow matching and Markov bridge frameworks. Furthermore, we extend our model to also explore the conformational landscape of the protein by jointly sampling side chain angles and molecules.
INTRODUCTION
Small molecules the predominant class of FIDA-approved drugs with share of 85%, and more than 95% of known drugs target human or pathogen proteins (Santos et al. 2017). At the same time, the cost and duration of the development of new drugs are skyrocketing (Simoens Huys [021). This sparks increasing interest in the computational design of small molecular compounds that bind specifically to disease-associated proteins and thus reduce the amount of costly experimental testing.
In recent years, the machine learning community has contributed plethora of generative tools ad- dressing drug design from various angles (Du al,, 2024). Some methods directly optimize specific drug properties, using techniques such as reinforcement learning (Popova et al. [018 Gottipati al., 2020) or search-based approaches
Hoomerzlommarellll et al. [018]] [Saanson et al.) (024))
However, these methods typically require careful tuming of the objective function to avoid exploiting imperfect computational oracles and overly maximizing one desired property (e.g. binding affinity)
at the expense of another (e.g. oral bioavailability). Additionally, 3D binding pose along with the chemical structure of the molecule, which substantially increases the degrees of freedom. Many optimization algorithms struggle efficiently navigate such vast design spaces.
often aims to design a suitable
Following different approach, probabilistic generative models learn generate drug-like molecules directly from data (Hoogeboom et al.] [022 Vignac et al. 2022). Here, the design objec- tives are implicitly encoded in the training data set. While these methods may not outperform direct optimization on isolated metrics, they are well suited for the multifaceted nature of drug design as they learn "what drug looks like" in more general way. Once trained on sufficient high-quality data, these models can capture a more holistic picture of the molecular space compared to models optimized for limited set of target metrics.
The strength of generative modeling lies in its ability to reproduce patterns seen in the training data. However, many prior works on generative molecular design have focused on absolute metric values in baseline comparisons (e.g. identifying the best docking scores synthetic accessibility esti- mates), even though the proposed models did not directly optimize these quantities. For example, Table 6 from the recent survey on generative structure-based drug design methods by Zhang et al. 2023b) highlights a common evaluation strategy used in this context. We argue that the evaluation
"These authors contributed equally

---

Published as a conference paper at ICLR 2025
h.e
gillcs fe), gettee (e)
In_
Figure I: Method overview. DRUGFLOW operates on continuous ligand atom coordinates dr, dis- crete atom types hu, and bond types er. Its extension FLEXFLOW additionally operates on continu- ous side chain angles Xt. As shown on the left, coordinates and angles are sampled from a Gaussian prior in 3D and uniform prior on the torus, respectively. Discrete types sampled from categor- ical prior distributions (uniform or marginal). Denoising schemes for continuous and discrete data types are based conditional flow matching (top) and Markov bridge models (bottom). As shown in the middle, at time step f, we process the noisy data to obtain the 3D graph: side chain angles are transformed into vector features pointing atom positions using NERF, and three types of edges are introduced using different distance cutoffs-edges between ligand atoms et"L edges between ligand and protein C. atoms e!*l, and edges between protein C. atoms PoP. The graph processed by neural network which outputs velocities (r t). (x(xe.t) t) and type predictions heg é, which are fed into the generative modeling framework. The model can adapt the size of the molecule during sampling and label excessive atoms that will be eventually removed (shown as crosses on the right). Additionally, DRUGFLOW outputs per-atom uncertainty values ōw.
of generative models for drug design should instead be centered around their ability to represent the training distribution accurately analogous to how image generation methods are typically as- sessed (Heusel al. 2017). Without additional fine-tuning or sampling strategies, it unreasonable to expect model to substantially improve any score compared to the training set.
In this work, we present DRUGFLOW, a new generative model for structure- sarueture-based drug design, that simultaneously learns the distribution of protein- binding molecules in three data domains. DRUGFLOW generates discrete atom and bond types as well atom coordinates in the Euclidean space. Its extended version, FLEXFLOW, additionally samples side chain configurations of the bind- ing pocket represented as angles on a hypertorus. This allows us to sample probabilistic ensembles of possible binding modes and enables drug design for targets in unbound conformations. Both DRUGFLOW and FLEXFLOW are conditioned on fixed protein backbone coordinates and amino acid types, which used as context for denoising. We further introduce virtual node type to allow the model to dynamically add or remove atoms and thus learn about the distribution of ligand sizes rather than requiring it be pre-specified. Finally, we add an uncertainty head to our model that is trained in an end-to-end fashion to identify out-of-distribution samples and rank molecules at inference time.
Based on our observations, we focus our evaluation primarily on the distribution learning capa- bilities of the proposed generative model, comparing assess molccular properties and structural features using distance functions between distributions de- rived from generated samples and training data points, and demonstrate that DRUGFLOW molecules closely match the data distribution across broad range of metrics. This suggests that DRUGFLOW can retrained on curated datasets to steer the generation of samples towards desired regions of the chemical space for various practical applications. However, recognizing that excessive filtering may lead to insufficient training data, we present an alignment strategy which allows us to update a pre-trained model based on user preferences.
established baselines. To this end, we
We summarize the main contributions of this work as follows.
Conceptual novelty We present DRUGFLOW and FLEXFLOW, new generative models for structure-basee drug design, and introduce three conceptually new features: (1) and end-to-end trained uncertainty estimate that successfully detects out-of-distributie samples, (2) an adaptive size selection method that discards excessive atoms during sampling, and (3) a protein conformation sampling module that samples realistic side chain rotamers.
2

---

Published as a conference paper at ICLR 2025
Performance We propose to evaluate generative models in way that better captures their training objective, and use the new benchmarking framework to show that DRUGFLOW is a state-of-the-an art distribution learner for struttue-basek drug design.
Practical relevance Recognizing that medicinal chemists ultimately aim to optimize molecules for specific design objectives, implement preference alignment scheme that allows us efficiently sample molecules with improved target properties. Recognizing that medicinal chemists ultimately aim to optimize molecules for specific design objectives, we implement preference alignment scheme that allows us to efficiently sample molecules with improved target properties.
METHODS
Our base model probabilistic model operating simultaneously on atom types, bond types and coordinates, thereby combining generative processes for discrete and continuous data types. We use Euclidean flow matching (Lipman et al., 2022) for ligand coordinates and combine with Markov bridge models (Igashov et al., (202)) applied to atom and bond types to generate the dis- crete molecular graphs. More background on these generative modeling frameworks is presented in Appendix A.1
The backbone model is heterogeneous graph neural network that has independent trainable weights for ligand and protein node types, as well as ligand, protein, and interaction edge types. Each protein node represents a whole residue but we include vector-valued input features to encode the locations of all atoms belonging to the residue. This allows to reduce computational complexity while preserving full atomic detail. The neural network predicts several node-level and edge-level outputs required for sampling, including logits for atom and bond types and vectors for the coordinates. All trainable operations applied to geometric quantities are implemented as geometric vector percep- trons (GVPs) (Jing et al. 2020) to ensure that predicted vector fields transform equivariantly. Full architectural details are provided in Appendix A.5
Below, we describe our conceptual novelties one by one, starting with uncertainty estimation in Section 2.1 followed by ligand size adaptation in Section (tttion2.221amm [2.2]and and protein flexibility in Section 2.3 Finally, we present a fine-tuning approach to align the general distribution learner with user specified preferences in Section 2.4 Our method is schematically depicted in Figure I
2.1 UNCERTAINTY ESTIMATION
To identify outtof-distribution (OOD) samples and endow DRUGFLOW with an intrinsic uncertainty estimate, we rely on technique that has been successfully used for regression problems in the past (Nix & Weigend, 1994 Lakshminarayanan
al.) 2017). We assume that the flow matching regression error is normally distributed with standard deviation dg, and derive a loss function that maximises the likelihood of the true vector field under this uncertainty model,
1
LFM.OOD F (.. log ( () 20 (x,, t)
[avo(a () 2.,1 2 (2), 112, (1)
where vx) E Rdd and co(2,,t) t) IR two output heads of the neural network and a; is the ground-truth conditional vector field. The derivation is provided in Appendix A.2
A model trained in this way provides us with per-atom uncertainty score in addition the vec- tor field for flow matching at every sampling step. We integrate the step-wise score to obtain an uncertainty value for the entire sampling trajectory (motivated in Appendix A.3) as
O - VV (z t)dt,
(2)
and assign this value to the resulting generated atom.
2.2 END-TO-END SIZE ESTIMATION
It is common practice to choose molecule sizes for diffusion and flow matching models in structure- based drug design a priori. The number of generated atoms is often either a user-specified hyper- parameter, sampled from the empirical distribution of sizes in the training set (Hoogeboom al.
3

---

Published as a conference paper at ICLR 2025
Schneuing et al. 2022 Guan et al. 2023a), estimated by a separate neural network (Igashov] et al. 2022). This approach, however, prevents the generative neural network from adapting to the context as the current sample evolves. For instance. if initially too many atoms specified it might be impossible for the model to create a molecule without steric clashes with the surrounding protein atoms.
2022
In order to also learn this aspect of the data distribution in an end-to-end manner, we aim to adapt the molecule size during the generative process, which effectively changes the dimension of the modeled system (Campbell et al. 2024). To do so, we introduce a virtual ("no atom") node type the model can sample. All nodes of this type are treated as completely disconnected (i.e. all bonds of virtual nodes have "None" type) and will be removed at the end of sampling. During training, order to treat virtual nodes and
we add Vivin ~ (0, Nmax) virtual nodes to each training sample. atoms in the same way. we need to attach coordinates to them as well. While different design choices are possible, we find that placing them the center of mass of the ligand works well in practice as it provides a clear reference point for the network to regress toward. Note that this approach still requires pre-specifying the number of nodes in the computational graph which serves as an upper bound for the number of atoms in the generated molecule.
2.3 PROTEIN FLEXIBILITY
Incorporating the dynamics and flexibility of protein structures is one of the key open challenges for structure-based drug design (Fraser & Murcko, 2024). As first step addressing scenarios in which the bound structure of the target protein is unknown or assuming a single static pocket configuration too restrictive, we extend our base model also generate side chain torsion angles for all pocket residues while sampling new ligands, thereby permitting full side chain flexibility.
To this end, we apply flow matching on the hypertorus that describes all torsion angles. Following the Riemannian Flow Matching framework Chen & Lipman (2023), we approximate vector fields in the tangent space. More details and definitions for flow matching on Riemannian manifolds is provided in Appendix A Similar to related works on flow matching for angular domains (Yim 2024) we find a non-linear scheduler n:(t), which controls how quickly the
et al., |023a|Lee geodesic distance between start and end point of trajectories decreases, to be beneficial to sample
quality. However, unlike these works we adopt a polynomial scheduler x(t) = (1 11-68 +)* with A: = 3. This has similar effect on the sampling trajectories as the exponential scheduler x(t) = e- if
c = 5 but strietly fulfills the theoretical requirement on the boundary conditions: s(0) = 1 and <(1) = O. Using the exponential and logarithm maps associated with the manifold, we derive the following updated flow and vector fields:
Kt exP no ((1 (1 t)at log (a1)),
(3)
c. A(1 1) k- logaro (211)
(4)
We use this scheduler both for training and sampling.
While the generative process performed entirely angular space to enforce physical plausibility, we present the full-atomic information more explicitly to the neural network and convert the side chain dihedral angles back to atom positions in every training and sampling step. This operation is performed efficiently using the Natural Extension Reference Frame (NERF) algorithm (Parsons)
et al. 2005 Alcaide et al. 2022) in parallel for each residue.
2.4
MULTI-DOMAIN PREFERENCE ALIGNMENT
Many real-world applications require generating molecules with properties underrepresented in the initial training data. To address this need, we employ an alignment scheme inspired by Direct Pref- erence Optimization (DPO), a technique originally used to align large language models (Rafailov) et al. 2023) and later adapted to diffusion models (Wallace et al. 2024).
Using a pre-trained reference model 0, we generate a synthetic dataset of preference pairs D {(x) }.).. where the winning samples x have more desirable molecular properties compared to the losing samples a (see Section for details on our preference datasets). To align our method with these preferences, we introduce another (aligned) model y, initialised with the weights of the
4

---

Published as a conference paper at ICLR 2025
Table 1: Wasserstein distance between marginal distributions of continuous molecular data (bond distances and angles), drug- liggliikeness (QED), synthetic accessibility (SA), lipophilicity (logP) and numbers of rotatable bonds (RB). The last column reports the ShaennShannon divergence between the joint distributions of four molecular properties (QED, SA, logP and Vina efficiency score). The best result highlighted in bold, the second best is underlined.
Top-3 bond distances
Top-3 bond angles
Molecular properties
Method
C-C
C-N
C=C
C-C=C
C-C-C
C-C-O
QED
SA
logP
RB
JSDai
POCKET2MOL
0.050
0.024
0.045
2.173
2.936
3,938
0.072
0.576
1.209
2.861
0.223
DIFFSBDD
0.041
0.039
0.042
3.632
8.166
7.756
0.065
1.570
0.774
0.928
0.274
TARGETDIFT
0.017
0.019
0.028
4.281
3,422
4.125
0,050
1.518
0.489
0.354
0.242
DRUGFLOW
0.017
0.016
0.016
0.952
2.269
1.941
0.014
0.317
0.665
0.144
0.099
Table 2: Wasserstein distance between distributions binding efficiency scores and normalized numbers of different protein-ligand interactions. The best result is highlighted in bold, the second best underlined.
Binding efficiency
Protein ligand interactions
Method
Vina
Gnina
H-bond
H- H-bond (acc.)
H-bond (don.)
r-stacking
Hydrophobic
2oL
0.064
0.044
0.040
0.026
0.014
0.007
0.027
DIFFSBDD
0.086
0.043
0.047
0.030
0.017
0.011
0.044
TARGETDIFF 0.034
0,030
0.031
0.021
0.010
0.012
0.039
DRUGFLOW 0.028
0.013
0.019
0.012
0.007
0.006
0.036
reference model 0. From now the reference model @ is no longer optimised, and further training is only performed on the new model y.
{coord, atom, bondj, we compute the loss terms (v) := La sp) and (v() == Ca(x', :p) for winning and losing samples, respectively. More specifically, these are the flow matching loss for coordinates (Eq. [2) and the Markov bridge model loss for atom and bond types (Eq. 24 Additionally, we calculate the corresponding loss terms (0) and ()) using the reference model 0 with fixed parameters. Using individually weighted loss differences AW L (5) - C (f) and A = (((()) (0), we define the multi-domain preference alignment (MDPA) loss follows,
For each data domain
( (MDPA(P) = - log a (-3. I. - -!.) + (x)
(5)
where Ae. Aw, A are adjustable weights, and a is the sigmoid function. Note that we regularize training by adding the overall loss terms L (1p) and C'((p) (Eq. (25) for the winning and losing samples, respectively. We find that such regularization significantly enhances training stability. More details are provided in Appendix A.4
EXPERIMENTS
3.1 MULTI-DOMAIN DISTRIBUTION LEARNING
In this work, we focus on the generative capabilities of our model. namely on its ability to learn the training data distribution. While a common trend in the community is to report absolute values of various molecular properties and docking scores, we stress that such an evaluation relevant only for methods whose primary goal is property optimisation, such preference alignment (Cheng cet al., 2024) or optimisation in the latent space
et al.) 2018). Unless stated specifically, DRUGFLOW does not optimise for any specific property and aims to learn the data distribution only. Therefore, instead of absolute values molecular properties we measure the proximity of distributions of these properties computed the generated samples and the training data. As we show in Section[3.5] absolute values of the target metrics can increased by fine-tuning on relevant data or using the preference alignment scheme.
5

---

Published as a conference paper at ICLR 2025
Metrics We compute bond types and ring systems (Walters) 2022 2021). We use the Wasserstein- distance for the bond length distributions of the three most common bond types (C-C. C-N and C=C), the three most common bond angles (C-C=C, C-C-C and C-C-O) as well as the number of rotatable bonds per molecule. We also apply the Wasserstein distance computational scores relevant applications in medicinal chemistry: Quantitative Estimate of Drug-likeness (QED) (Bickerton et al., 2012). Synthetic Accessibility (SA) (Ertl & Schuffenhauer 2009) and lipophilicity (logP) (Wildman & Crippen 1999). Binding efficiency, defined as a computational binding affinity score divided by assessed in the same manner. We report efficiency scores instead of affinity scores due to the high correlation of the latter with molecule size (Cremer et al.) 2024). We use both Vina and Gnina docking scores (McNutt et al. 202T) as binding affnnity oracles. Besides, we compare normalized counts of various types of non-covalent interactions as detected by ProLIF (Bouysset & Fiorucci 202I). To do this, we divide the number of interactions by the number of atoms in each molecule and compute Wasserstein distances to compare the resulting distributions.
divergences for the categorical distributions of atom types,
the number of atoms in the molecule,
Finally, we also employ the Fréchet ChemNet Distance (FCD) (Preuer et al.) 2018 to assess how well the model approximates the training data distribution.
Dataset & Baselines We use the CrossDocked dataset (Francoeur et al. 2020) with 100 000 protein-ligand pairs for training and 100 proteins for testing, following previous works (Luo al. 2021 Peng et al. 2022). The data split was done by 30% sequence identity using MMseqs2 (Steinegger & Söding 2017). Ligands that do not pass all J
et al. (2024) fil- ters were removed from the training set. We compare DRUGFLOW with an autoregressive method, 2OCKET2MOL (Peng et al. 2022). and two diffusion-based methods, TARGET (Guan et al.) 2023a) and DID (Schneuing et al. (022). We generated 100 samples for each test set target with DRUGFLOW and selected only molecules that passed the RDKit validity filter.
Results The distances between sampling dis- tributions measured on the various discrete and continuous characteristics of molecules and in- teractions are summarized in Tables [ [2] and 3 DRUGFLOW shows convincing all-round performance and outperforms other methods in almost all aspects. POCKET2MOL achieves slightly better results on bond types and hy- drophobic interactions. TARGETDIFF gets the best results on the lipophilicity metric. How- ever, in these few cases DRUGFLOW consis- tently ranks close second. Along with the marginal distributions of various geometric and themo-physicaa characteristics, we compare joint distributions of the molecular properties. To do this, we created histograms of the joint distributions of QED, SA, logP and Vina efficiency scores (10 bins per score, Figure and compute the Jensen-Shannon divergence between them. DRUGFLOW outperforms other methods by a large margin, as demonstrated by JSD in Table D Addition- ally, our method substantially outperforms other baselines in FCD. To further analyse this result, we applied Principal Component Analysis (PCA) to reduce the dimensionality of the ChemNet em- beddings and visualized the first two principal components on a 21 plane. As shown Appendix Figure DRUGFI.OW covers considerably more modes of the training distribution compared to the other methods. Distributions of other metrics are provided in Appendix B.T Additionally, we com- pute novelty, uniqueness, and overall quality of the samples using PoseBusters (Buttenschoen al.) 2024]. As shown in Appendix Table 6 DRUGFLOW demonstrates competitive results in all these metrics well. Appendix Tables 7 g and 9 provide additional results for different variations of DRUGFLOW and FLEXFLOW. The statistical significance is ensured in Appendix B.
Table 3: Fréchet ChemNet Distance and Jensen- Shannon divergence between distributions of dis- crete molecular data. The best result is high- lighted in bold, the second best is underlined.
Method
FCD
Atoms
Bonds
Rings
POCKET2MOL
12.703
0.081
0.044
0.446
DIFSSDD
11.637
0.050
0.227
0.588
TARGETDIFF
13.766
0.076
0.240
0.632
DRUGFLOW
4.278
0.043
0.060
0.391
3.2 OUT-OF-DISTRRIUTTION DETECTION
The uncertainty estimation technique described in Section 2.1 provides per-atom uncertainty scores, as shown in Figure A Here, we showcase six molecules of the same size with the highest and lowest global (L.e. averaged over atoms) uncertainty scores.
To demonstrate the ability of our uncertainty score to detect outtt--diitributiin samples, we plotted histograms of several basic geometric properties of the generated molecules. The histogram bins

---

Published as a conference paper at ICLR 2025
A
C
Uncertainty of clashing and non-clashing atoms
High uncertainty
1
073
Gnina eff,
Uncertainty
0.90
0.90
0.9
Low uncertainty
0.58
0.8
T- 0.88
0.84
Clashes No clashes
B
Uncertaintyy
C-C-C
C-C-O
0.46
120 Bond angles,
0.800
Bond angles, deg.
Bond angles, deg.
Bond length,
A
Bond length,
Bond length, A
Figure 2: Uncertainty estimate detects out-of-distribution samples and correlates with various molecule characteristics. (A) Examples of samples with the same size (30 heavy atoms) and high and low global uncertainties. Each atom highlighted according to its local uncertainty score. (B) Distributions of top-3 bond angles and bond lengths in DRUGFLOW samples (histograms) and training data (dashed lines). The histogram bins are color-coded according to the average uncertainty values of the corresponding data points. In all cases, samples from the distribution tails have high uncertainty. Main modes of the distributions have lower uncertainty on average. The color bar is scaled for a better visibility, (C) Correlation of the global uncertainty score with Gnina efficiency score and size of the molecules. (D) Atoms clashing with the protein have higher uncertainty scores than non-clashing atoms.
color-coded according to the average uncertainty values of the corresponding data points. As shown in Figure [23, the model tends to assign high uncertainty to data points in the tails of the distributions, while sampling from the central regions (modes) with high confidence. The slight limitation is that most of the values fall within the range of 0.85 to 0.92. However this does not affect the discriminative ability of the score.
Furthermore, our analysis reveals strong correlations between global uncertainty and both molecule size and docking efficiency, shown in Figure C While the correlation with molecule size is expected, the correlation with size-agnostic ligand efficiency scores is particularly interesting and has practical implications. We hypothesize that this correlation may distribution learning process, wherein, among other data aspects, the model learns to avoid steric clashes between protein and ligand atoms. Indeed, as shown in Figure 2D, atoms sampled without clashes tend to have lower uncertainty values compared to clashing atoms.
byproduct of efficient
3.3 DISTRIBUTION OF MOLECULES SIZES
As explained in Section 2.2 every training step add the molecule random number of virtual nodes sampled from (0. Nmax). It means that on average the model was trained to label as virtual (i.e. remove) N Mmxx/ /2 atoms per molecule. In our experiments, we set Nmax 10, and therefore expect the model to remove on average 5 atoms during sampling. To test this hypothesis, we sampled molecules with ground-truit sizes and 5 atoms added on top. As shown in Figure the model indeed tends to remove about 5 atoms. The deviation from the uniform training distribution (shown in red) is evidence of the model's ability to learn the conditional distribution of molecule sizes given pockets. CrokeDocked pockets are defined based distance cutoff, which introduces a dependency between ligand and pocket sizes.
Next, studied the ability of the model to minimise steric clashes with the pocket when large number of (computational) nodes is provided. To maximise the number of geometric constraints, we selected protein with deeply buried pocket (PDB: II.3I.) and tightly bound ligand (17 atoms),
7

---

Published as a conference paper at ICLR 2025
Clashes within the learned "bandwidth"
Clashes beyond the learned "bandwidth"
Virtual (removed)
A
B
nodes in samples
C
Distributions of virtual nodes
clashing atoms
wio virtual nodes w/ vittual modes
removed nodes
14
during training In samples
12
40-
10-
30-
20
60
Input molecule size
Input molecule size
Input molecule size
Number of virtual nodes
Figure 3: DRUGFLLOW learns the conditional size distribution of molecules given protein pock- ets. (A) Our model effectively removes redundant atoms to avoid clashes. (B) Trained with maxi- mum Maxx 10 virtual nodes, the model struggles remove much more atoms. (C) Even though the training distribution of number of added virtual nodes is (0, Nmax), the distribution of the re- moved nodes during sampling suggests that the model learned the dependency between ligand and pocket sizes.
as shown in Appendix Figure 16 For different input sizes (between 17 and 67 atoms), we sampled 1000 molecules per each size value and measured the percentage of clashing atoms and numbers of atoms removed by the model. We repeated the same experiment with similar model trained without virtual nodes. As shown in Figure 33, the model produces minimum clashes long the number of excessive atoms does not exceed 10, the "bandwidth" the model was trained with. Higher numbers out of the training distribution, and therefore the model fails to remove more, as shown in Figure 3 We helieve that scaling up the maximum number of virtual nodes during training will enable the model operate in fully adaptive size selection regime.
3.4 LEARNING THE DISTRIBUTION OF SIDE CHAIN ROTAMERS
Here evaluate how well FLEXFLOW, the version of our model that simultancously generates side chain conformations, recovers the distribution bound side chain rotamers. Because all proteins in test are provided in ligand-bound form, repacked their side chains in absence of the small molecule to approximate their unbound structures. We relaxed side chains with the Rosetta repack protocol (Conway et al. 2014) and achieved RMSD of about 1.98A for side chain heavy atoms (Figure 4A, middle).
Recovering bound configurations As first test, we sampled 20 sets of side chain torsion angles per test set target with FLEXFLOW while keeping the ligand fixed. We achieve this by using the ground truth vector field and transition probabilities instead of the predicted quantities for all ligand- related variables. Figure 4 shows that the model samples pocket structures close to the original bound conformations (with median side chain of RMSD 1.75Ā). This decrease in RMSD compared to the unbound pocket is expected because fixing the ligand binding pose constrains the space of feasible solutions. For comparison, we also include the distribution of RMSD values that results from simply taking random angles from the prior distribution (Figure left).
Distribution of side chain angles Next, we freely generate molecules and bound configurations of the protein-ligand complex and assess whether the resulting side chain rotamers are in accor- dance with the bound structures from the training set. Starting from a completely random prior, FLEXFLOw manages to recover the rotameric modes of the reference structures accurately. Fig- ure 4 shows the distributions of the first two side chain torsion angles three bulky amino acids. Analogous plots for all 14 amino acids with at least two side chain angles are presented in Appendix Figure
3.5 PROPERTY OPTIMIZATION WITH PREFERENCE ALIGNMENT
We conduct preference alignment with respect to four molecular properties: QED, SA, Vina effi- ciency score, and Rapid Elimination of Swill (REOS) filters (Walters et al.) 1998). REOS filters in- clude various structural alerts designed to detect problematic compounds in screening libraries (Baell
& Holloway 2010 Hann al. 1999 Pearce al. 2006). For REOS, molecules classified as
8

---

Published as a conference paper at ICLR 2025
Samples
Training set
A
B
ARG
LYS
TRP
RMSD (A)
3
(rad)
2
1
Prior Unbound
Generated
-2
0 (rad)
2
4
4
-2
4
4
-2
0 (rad)
2
4
(rad)
Figure 4: FLEXFLOW samples realistic side chain conformations. (A) Side chain root-mean- square deviation (RMSD) for random samples from the model's prior (left), relaxed pockets in absence of the ligand (middle), and FLEXFLOO-generae structure (right). (B) Distributions of XI and X2 angles for Arginine, Lysine and Tryptophan. We compare FLEXFL.OW samples to the bound pocket conformations from the training set.
side chain conformers for fixed ligand
either passing all REOS filters (winning) or failing at least one (losing). For continuous metrics, we set thresholds requiring the winning sample to outperform the losing one by at least 0.5 for SA, and 0.1 for both QED and Vina efficiency. In each pair, both molecules are generated for the same pocket. In addition, curate dataset where winning samples surpass losing samples across all four properties. To perform preference alignment, we initialize the model with the reference DRUGFL.OW model parameters and train until convergence of the loss LMDPA (Eq. (5). We then select checkpoint using the validity metric on the validation set. To compare our preference align- ment method against a simpler optimization strategy, we fine-tune the same reference model on the winning samples only.
Figure shows the performance gains of the aligned and fine-tuned models compared to the training data and the reference model. Additionally, we visualize distribution shifts of the target metrics in Appendix Figure [8] and show their statistical significance the preference-aligned models consistently exceed their fine-tuned counterparts across all metrics. However, these improvements are achieved at the cost of moderately reduced molecular validity (Appendix Table (3). Notably, the model optimised for all four properties at once demonstrates competitive results across all target metries.
Figure 19 Our results indicate that
RELATED WORK
Generative models for molecule generation This paper builds on large body of work on proba- bilistic models for molecule generation. Some of these models generate molecules unconditionally without knowledge of the target protein. They either create molecular graphs (Jo et al. 2022 Vignac
et al. 2022). 3D atomic point clouds (Gebauer al. 2019 Gareia Satorras et al. 2021 Hoogeboom ct al. 2022) or both Vignac et al. (002)). More closely related to the work presented here, another family of models attempts to generate novel chemical matter conditioned
dimee-dimenion
text, typically a structural model of a target protein. Among these, most of the models sample atom
positions and types without providing explicit information about covalent bonds (Liu et al. 2022
Ragoza et al. [022]] Schneuing et al., 2022, [gashov et al. 2022 Guan et al. 2023a Lin et al.
2022
Xu al. 2023). Others generate the full molecular graph structure and binding pose jointly (Peng
et al. 2022]] Guan et al., 2023b Zhang et al. 2023a). Notably, most recent 3D models belong to
the family of diffusion probabilistic models (Schneuing et al.) 2022 Guan et al., 2023a| Lin al.)
2023 Guan et al. 2023b Weiss et al. (223)) or follow the related flow matching paradigm (Song al.) 2024; Dunn Koes 2024) Irwin al. (024). So far, it is uncommon for
2022
Xu et al.
these models to handle varying dimensionality (i.c. molecule sizes) (Campbell et al. 2024) or incor- porate protein flexibility although torsional flow matching has already been applied to protein side
chain packing (Lee & Kim 2024) and peptide design (Lin al. 2024
Confidence prediction Neural confidence estimates are crucial components of many popular meth- ods for biomolecular applications (Jumper al. 2021 Corso al. 2022 Abramson et al. 2024). Typically, such estimates are obtained by training a separate neural network or auxiliary output head to approximate the prediction error of the main model either during training (Jumper et al. 2021

---

Published as a conference paper at ICLR 2025
REOS
QED
Vina efficiency
SA
1.0
0.89
0.9
0.11
0.0
0.75
0.66
0.5
0.53
0,58
-0.2
3.22
3.47
3-
0.4
0.4
-333
2.39
-0.4
0.3
0.40
0.2
0.0
0.0
Tramiin set
Reference
Fine tndd
PA (candu
Peel. align.
Raining set
Reference
Fine tureed
PA (comb.
Pren alggn
Training set
Reference
Fine Fine-tuned
PA (cond.)
Pref. aign.
Rorerencc
Fine mee-uuee
Pk (contu)
Pet algn.
Training
Figure 5: Preference alignment improves target properties. Comparison of models across four molecular properties: REOS, QED, SA, and Vina efficiency. Bars represent the training set (gray), reference model (green), fine-tune munet model (blue), preference-aligned model with combined pref- erences (dashed), and preference-aligueee Preference-aligee aligned models consistently outperform fine-tuning, with the best results in the metric- specific models. Results are based on 100 sampled molecules per target.
model for the metric specified above each plot (red).
Abramson et al. 2024) or for samples from a trained model (Corso et al. 2022 Unlike our ap- proach, this requires minimizing a separate secondary loss function. Furthermore, to enable this form of confidence estimation based final outputs during training, diffusion (or flow matching) models like AlphaFold 3 (Abramson et al., (202)) must perform expensive rollout schemes to regress their output errors directly. Here, circumvent this by deriving a global uncertainty score from local, step-wise predictions.
Preference alignment Generating molecules with desired properties is critical step in applying generative models to biomolecules. Existing methods, such as reinforcement learning (RL)-based approaches (Zhou et al. 2019) GANs (Maziarka al., 2020), often involve two-step process: first. to generate molecules, then optimize them. In contrast, preference alignment techniques allow generative models directly produce outputs aligned with human or domain- speccifc preferences in single step. DPO offers key advantage over RL by eliminating the need to learn a reward function, making it a less exploitable way to optimize preferences (Rafailov et al., 2023). DPO has been applied to fine-tune chemical language models (Park et al., (2033) and its adaptation for diffusion models (Wallace et al. 2024) has already found applications in antibody design (Zhou et al. 2024) and structure-based drug design (Cheng et al. 2024 Gu et al. (024). In DecompDPO (Cheng al. 2024]. molecules are decomposed into fragments, allowing preference optimization at both local and global levels, which helps resolve conflicting preferences between molecular substructures. In AliDiff (Gu et al. 2024). the DPO objective decomposed across atom types and coordinates and rescaled according reward differences in preference pairs. Our approach extends this by applying a preference alignment scheme simultaneously in different generative frameworks (continuous flow matching and discrete Markov bridge models).
5 CONCLUSION
In this work, we advocate for a distribution learning-centered evaluation of generative models for drug design. Methods that perform well in this framework, can later improve their task-specific performance through curated datasets, advanced sampling strategies, or fine-tuning to better align with user preferences. We introduce DRUGFLOW and demonstrate that it consistently achieves state- of-the-art distribution learning performance across various orthogonal metrics. DRUGFLOW also learns the distribution of molecular sizes and is able to detect out-of-distribution samples. Besides, its extension, FLEXFLOW, additionally operates on protein residues and learns the distribution of side chain conformations. Finally, discuss a preference alignment strategy that allows to sample molecules with improved properties, and show its effectiveness across four different metrics.
REPRODUCIBILITY STATEMENT
All methodological details are described in Section and Appendix https: //gi http://ggggtuu.......ooooooooo -ttp:LLLPOT-PPF//ruggg ow
Source code available at
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors would like to thank Rebecca Neeser, Julian Cremer, Charles Harris, Felix Dammann, Yuanqi Du, Evgenia Elizarova, and Simon Crouzet for their helpful feedback. This work was sup- ported by the Swiss National Science Foundation grant 110030_197722 I.I. has received funding from the European Union's Horizon 2020 Research and Innovation Programme under the Marie Aklodowska-Curte grant agreement No. 945363. A.W.D. was supported by the Studienstiftung des deutschen Volkes. M.B. is partially supported by the EPSRC Turing AI World-Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub No. EP/Y028872/1.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic inter- polants, arXiv preprint arXiv:2209. 2022.
Eric Alcaide, Stella Biderman, Amalio Telenti, and M Cyrus Maher. Mp-nerf: A massively parallel method for accelerating protein structure reconstruction from internal coordinates. Journal of Computational Chemistry. 43(1):74-78, 2022.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:1 7981-17993, 2021.
Jonathan Baell and Georgina A Holloway. New substructure filters for removal of pan assay inter- ference compounds (pains) from screening libraries and for their exclusion in bioassays. Journal of medicinal chemistry, 53(():2719
2010.
G Richard Bickerton, Gaia V Paolini, Jérémy Besnard, Sorel Muresan, and Andrew L Hopkins. Quantifying the chemical beauty of drugs. Nature chemistry, ((2):90-98, 2012.
Cédric Bouysset and Sébastien Fiorucci. Prolif: library to encode molecular interactions as finger- prints. Journal of cheminformatiis, 13(1):72, 2021.
Martin Buttenschoen, Garrett M Morris, and Charlotte Deane. Posebusters: Ai-based docking methods fail to generate physically valid poses or generalise novel sequences. Chemical Sci- ence, 2024.
Andrew Campbell, William Harvey, Christian Weilbach, Valentin De Bortoli, Thomas Rainforth, and Arnaud Doucet. Trans-dimensional generative modeling via jump diffusion models. Ad- vances Neural Information Processing Systems, 36, 2024.
Ricky TQ Chen and Yaron Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03660. 2023.
Ting Chen, Ruixiang Zhang, and Geoffrey Hinton. Analog bits: Generating discrete data using diffusion models with self-conditioning. arXiv preprint arXiv:2208.04200: 04202, 2022.
Xiwei Cheng, Xiangxin Zhou, Yuwei Yang. Yu Bao, and Quanquan Gu. Decomposed direct prefer- ence optimization for structure-based drug design. arXiv preprint arXiv:2407.13981 2024.
Patrick Conway, Michael D Tyka, Frank DiMaio, David E Konerding, and David Baker. Relaxation of backbone bond geometry improves protein energy landscape modeling. Protein science. 23(1): 47-55, 2014.
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Dif- fusion steps, twists, and turns for molecular docking, arXiv preprint 11X01:2210.0177,, 2022.
11

---

Published as a conference paper at ICLR 2025
Julian Cremer, Tuan Le, Frank Noé, Djork-Amé Clevert, and Kristof T Schütt. Pilot: Equivariant diffusion for pocket conditioned de novo ligand generation with oujti-objective guidance via importance sampling. arXiv preprint arXiv:2405 2024.
Yuanqi Du, Arian R Jamasb, Jeff Guo, Tianfan Fu, Charles Harris, Yingheng Wang, Chenru Duan, Pietro Liò, Philippe Schwaller, and Tom L Blundell. Machine learning- aided generative molecular design. Nature Machine Intelligence, pp. 1-16, 2024.
Ian Dunn and David Ryan Koes. Mixed continuous and categorical flow matching for 3d de novo molecule generation. arXiv preprint arXiv:2404.19733 2024.
Peter Ertl and Ansgar Schuffenhauer. Estimation of synthetic accessibility score of drug-like molecules based on molecular complexity and fragment contributions. Journal of cheminfor- matics, 1:1-11, 2009.
Paul G Francoeur, Tomohide Masuda, Jocelyn Sunseri, Andrew Jia, Richard B lovanisci, Ian Snyder, and David R Koes. Thressaimeniional convolutional neural networks and cross-docked data set for structure-based drug design. Journal of Chemical Information and Modeling, 60(9):4200- 4215, 2020.
James S Fraser and Mark A Murcko. Structure beauty, but not always truth. Cell. 18763)5177-500. 2024.
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E (n) equivariant normalizing flows. Advances in Neural Information Processing Systems, 34:4181- 4192, 2021.
Niklas Gebauer, Michael Gastegger, and Kristof Schütt. Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules. Advances in neural information processing systems, 32, 2019.
Rafael
Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sänchez-Lengeling. Dennis Sheberla, Jorge gguilera-Iparraguirre Timothy D Hirzel, Ryan P Adams, and Alán Aspuru-Guzit Automatic chemical design using data-driven contin- uous representation of molecules. ACS central science, 4(2):268 276, 2018.
Sai Krishna Gottipati, Boris Sattarov, Sufeng Niu, Yashaswi Pathak, Haoran Wei, Shengchao Liu, Simon Blackburn, Karam Thomas, Connor Coley, Jian Tang, et al. Learning to navigate the synthetically accessible chemical space using reinforcement learning. International conference on machine learning, pp. 3668-36679 PMLR, 2020.
Siyi Gu, Minkai Xu, Alexander Powers, Weili Nie, Tomas Geffner, Karsten Kreis, Jure Leskovec, Arash Vahdat, and Stefano Ermon. Aligning target-aware molecule diffusion models with exact energy optimization. arXiv preprint arXiv:2407.01644 2024.
Jiaqi Guan, Wesley Wei Qian, Xingang Peng, Yufeng Su, Jian Peng, and Jianzhu Ma. 3d equiv- ariant diffusion for target-aware molecule generation and affinity prediction. arXiv preprint
arXiv:2303.0354. 2023a.
Jiaqi Guan, Xiangxin Zhou, Yuwei Yang. Yu Bao, Jian Peng, Jianzhu Ma. Qiang Liu, Liang Wang, and Quanquan Gu. Decompdiff: Diffusson models with decomposed priors for structure-based drug design. 2023b.
Mike Hann, Brian Hudson, Xiao Lewell, Rob Lifely, Luke Miller, and Nigel Ramsden. Strate- gic pooling of compounds for high-throughhus screening. Journal of chemical information and computer sciences, 39(5):897 902, 1999.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3d. In International conference on machine learning, pp. 8867- 8887. PMLR, 2022.
12

---

Published as a conference paper at ICLR 2025
Ilia Igashov, Hannes Stärk, Clément Vignac, Victor Garcia Satorras, Pascal Frossard, Max Welling, Michael Bronstein, and Bruno Correia. Equivariant 3d-conditional diffusion models for molecular linker design. arXiv preprint arXiv:2210.0522 2022.
Ilia Igashov, Ame Schneuing, Marwin Segler, Michael Bronstein, and Bruno Correia. Retrobridge: Modeling retrosynthesis with markov bridges. arXiv preprint arXiv:2308. 162/2, 2023.
Ross Irwin, Alessandro Tibo, Jon-Paul Janet, and Simon Olsson. Efficient 3d molecular generation with flow matching and scale optimal transport. arXiv preprint arXiv:2406.07266. 2024.
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael JL Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. arXiv preprint urXiv:2009.0141, 2020.
Bowen Jing, Stephan Eismann, Pratham N Soni, and Ron O Dror. Equivariant graph neural networks for 3d macromolecular structure. arXiv preprint arXiv: (0ti::2006.034433 2021.
Jaehyeong Jo, Seul Lee, and Sung Ju Hwang. Score-based generative modeling of graphs via the system of stochastic differential equations. In International Conference on Machine Learning. pp. 10362-10383, PMLR, 2022.
John Jumper, Richard Evans, Alexander Pritzel. Tim Green, Michael Figurnov, Olaf Ronneberger,
Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-58)) 2021.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances neural information processing systems, 30, 2017.
Jin Sub Lee and Philip M Kim. Flowpacker: Protein side-chain packing with torsional flow match- ing. bioRxiv, pp. 2024-07, 2024.
Haitao Lin, Yufei Huang, Meng Liu, Xuanjing Li. Shuiwang Ji, and Stan Li. Diffbp: Generative diffusion of 3d molecules for target protein binding, arXiv preprint a1X11:2211.11214, 2022.
Haitao Lin, Odin Zhang, Huifeng Zhao, Dejun Jiang, Lirong Wu, Zicheng Liu, Yufei Huang, and Stan Z Li. Ppflow: Target-aware peptide design with torsional flow matching. bioRxiv, pp. 2024- 03, 2024.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hame Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Meng Liu, Youzhi Luo, Kanji Uchino, Koji Maruhashi, and Shuiwang Ji. Generating 3d molecules for target protein binding. arXiv preprint arXiv:2204.09410 2022.
Shitong Luo, Jiaqi Guan, Jianzhu Ma, and Jian Peng. A 3d generative model for structure-based drug design. Advances in Neural Information Processing Systems,
2021.
Lukasz Maziarka, Agnieszka Pocha, Jan Kaczmarczyk, Krzysztof Rataj, Tomasz Danel, and Michat Warchot. Mol-cyclegan: generative model for molecular optimization. Journal of Cheminfor- matics, 12(1):2, 2020.
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 learn- ing. Journal of cheminformatics, 13(1):1-20, 2021.
David Nix and Andreas Weigend. Estimating the mean and variance of the target probability distribution. In Proceedings 1994 ieee international conference on neural networks (ICNN'94), volume 1, pp. 55-60. IEEE, 1994.
Ryan Park, Ryan Theisen, Rayees Rahman, Anna Cichońska. Marcel Patek, and Navriti Sahni. Pref- erence optimization for molecular language models. In NewrIPS 2023 Generative Al and Biology (GenBio) Workshop, 2023. URL. htR.https: http:////peeeeeee.eeeeeee net.:/iwwwweeeeeeeemmmmmeeeeee formm att.:e////oooo.mtttttttataaaaar
13

---

Published as a conference paper at ICLR 2025
Jerod Parsons, J Bradley Holmes, J Maurice Rojas, Jerry Tsai, and Charlie EM Strauss. Practi- cal conversion from torsion space to cartesian space for in silico protein synthesis. Journal of computational chemistry, 26(10):1033-1066 2005.
Bradley C Pearce, Michael Sofia, Andrew Good, Dieter M Drexler, and David A Stock. An empirical process for the design of tightthroughptt screening deck filters. Journal of chemical information and modeling, 46(3): (063):060010060 2006.
Xingang Peng, Shitong Luo, Jiaqi Guan, Qi Xie, Jian Peng, and Jianzhu Ma. Pocket2mol: Effi- cient molecular sampling based on 3d protein pockets. In International Conference on Machine Learning, pp. 17644-17655. PMLR, 2022.
Mariya Popova, Olexandr Isayev, and Alexander Tropsha. Deep reinforcement learning for de novo drug design. Science advances, 4(7)xcaap7885, 2018.
Kristina Preuer, Philipp Renz, Thomas Unterthiner, Sepp Hochreiter, and Günter Klambauer. Fréchet chemnet distance: A metric for generative models for molecules in drug discovery. arXiv preprint /803. 08, 2018.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model secretly a reward model. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https:
//openreview.
d==iu IXJaa
Matthew Ragoza, Tomohide Masuda, and David Ryan Koes. Generating 3d molecules conditional on receptor binding sites with deep generative models. Chemical science, 13(9):2701-2713, 2022.
Rita Santos, Oleg Ursu, Anna Gaulton, A Patrícia Bento, Ramesh S Donadi, Cristian G Bologa, Anneli Karlsson, Bissan Al-Lazikani, Anne Hersey, Tudor I Oprea, et al. comprchensive map of molecular drug targets. Nature reviews Drug discovery, 16(1):19-34) 2017.
Arne Schneuing, Yuanqi Du, Charles Harris, Arian Jamasb, Ilia Igashov, Weitao Du, Tom Blundell, Pietro Lió, Carla Gomes, Max Welling, Michael Bronstein, and Bruno Correia. crructure-based drug design with equivariant diffusion models. arXiv preprint arXiv:2210. 2022.
Steven Simoens and Isabelle Huys. R&d costs of new medicines: A landscape analysis. Frontiers in Medicine, 8:760762, 2021.
Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao, Yanyan Lan, Stefano Ermon, Hao Zhou, and Ying Ma. Equivariant flow matching with hybrid probability transport for 3d molecule generation. Advances in Neural Information Processing Systems, 36, 2024.
Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Harmonic self-conditioned flow matching for multi-ligand docking and binding site design. arXiv preprint arX1v:2310.05764, 2023.
Martin Steinegger and Johannes Söding. MMseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature Biotechnology, 15(11)10026 -00000002200220000000102022122 October 2017. doi: 10.10388/bt.3988.
1038/ not 3988
Kyle Swanson, Gary Liu, Denise Catacutan, Autumn Arnold, James Zou, and Jonathan M Stokes. Generative ai for designing and validating easily synthesizable and structurally novel antibiotics. Nature Machine Intelligence, 6(3):338 -53, 2024.
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Improving and generalizing fow-based generative models with minibatch optimal transport. In ICML Workshop New Frontiers in Learning, Control, and Dynamical Systems, 2023.
Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pas- cal Frossard. Digress: Discrete denoising diffusion for graph generation. arXiv preprint arXiv:2209. 14734, 2022.
14

---

Published as a conference paper at ICLR 2025
Clement Vignac, Nagham Osman, Laura Toni, and Pascal Frossard. Midi: Mixed graph and 3d denoising diffusion for molecule generation. arXiv preprint arXiv:2302.09044, 2023.
Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8238, 2024.
Pat Walters. Useful rdkit utils, 2021. URL
rrdki__uiis
Mining ring systems in molecules for fun and profit, 2022. httppp//paactealaaaaaaaaaaaaaaca ctt./c22/122222 Lee.tu/.....ee--oo---yeeeeee--er
Pat Walters.
URL
html
W Patrick Walters, Matthew T Stahl, and Mark A Murcko. Virtual screening-an overview. Drug discovery today, 3(4):160-178, 1998.
Tomer Weiss, Eduardo Mayo Yanes, Sabyasachi Chakraborty, Luca Cosmo, Alex M Bronstein, and
Renana Gershoni-Porane Poranne. Guided diffusion for inverse molecular design. Nature Computational Science, pp. 1-10, 2023.
Scott A Wildman and Gordon M Crippen. Prediction of physicochemical parameters by atomic contributions. Journal of chemical information and computer sciences, 39(5):868-88) -873, 1999.
Minkai Xu, Alexander S Powers, Ron O Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. In International Conference on Machine Learning, pp. 38592-38610. PMLR, 2023.
Jason Yim, Andrew Campbell, Andrew YK Foong, Michael Gastegger, José Jiménez-Luna, Sarah Lewis, Victor Garcia Satorras, Bastiaan Veeling, Regina Barzilay, Tommi Jaakkola, al. Fast protein backbone generation with se (3) flow matching. arXiv preprint arXin:2310.05527, 2023a.
Jason Yim, Brian L. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprini ar arXiv:2302.022
2023b.
Odin Zhang, Jintu Zhang, Jieyu Jin, Xujun Zhang, RenLing Hu, Chao Shen, Hanqun Cao, Hongyan Du, Yu Kang, Yafeng Deng, et al. Resgen is a pocket-aware 3d molecular generation model based on parallel multiscale modelling. Nature Machine Intelligence, 5(9): 2023a.
Zaixi Zhang, Jiaxian Yan, Qi Liu, Enhong Chen, and Marinka Zitnik. systematic survey in geo- metric deep learning for structure-based drug design. arXiv preprint 1rxiv:2366.17763 2023b.
Xiangxin Zhou, Dongyu Xue, Ruizhe Chen, Zaixiang Zheng, Liang Wang, and Quanquan Gu. Antigen-specific antibody design via direct based preference optimization. arXiv preprint arXiv:2403.16576. 2024.
Zhenpeng Zhou, Steven Kearne,, Li Li, Richard N Zare, and Patrick Riley. Optimization of molecules via deep reinforcement learning. Scientific reports, 911):10752, 2019.
15

---

Published as a conference paper at ICLR 2025
Appendix for "Multi-domain Distribution Learning for De Novo Drug Design"
TABLE OF CONTENTS
Extended Methods
17
Generative framework
17
A.1.1 Flow matching
17
A.1.2
Riemannian conditional flow matching
17
A.1.3
Markov bridge model
18
A.1.4
Training loss
18
A.2
Predictive uncertainty estimates for regression problems
19
Motivation for the final uncertainty score
20
A.4
4.4
Preference alignment
21
A.5
Model architecture and training
Extended Results
24
B.1
Extended distribution learning metrics
24
B.
2
Absolute metries
................................
27
Ablation studies
28
B.4 4
Uncertainty estimation
................................
29
Virtual nodes
................................
30
B.6
6
Distribution of side chain angles
................................
32
B.7
Additional preference alignment results
B.8
Prediction error
34
B.9
Distribution the joint QED, SA, logP and Vina scores
35
B.10
Significance of Differences in Wasserstein Distance and Jensen- Shannon Divergence
35
16

---

Published as a conference paper at ICLR 2025
A EXTENDED METHODS
A.1 GENERATIVE FRAMEWORK
A.1.1 FLOW MATCHING
Flow matching describes class of deep generative models that approximate deeeddepnddntt vector field (x), x: E IR which generates sequence of probability distributions (px 1.1 E
[0, 1]}
pushing prior p = po towards the data distribution g = P1- The flow the [0, 1] x Rdd R is
defined through an ordinary differential equation (ODE) given the vector field u,;
d (x) = u. (,(1 (a)). dt
(6)
Efficient training of flow matching models is only possible because does not need to define the true vector field (x( but can instead match the conditional flow u, (2|æz) which is much easier to parameterize (Lipman et al. 2022) based data point ø1. Thus, the conditional flow matching loss amounts to
LCFM (0)
(x.g((1).p vu( (r 1) -)=
where ã; d/dt ( 1 is the time derivative of the conditional flow æt = the (
For sampling, we obtain sample zu from the prior p and simulate the ODE in Eq. replacing the true vector field u, with the learned vector field vo (x,,,))
In this work, we build on variant called Independent-coupling Conditional Flow Matching (ICFM) (Albergo & Vanden-Eljnnen 2022, Tong et al. 2023) and consider Gaussian conditional
probability path
We use this setup with model the flow for ligand coordinates. This results in constant velocity vector field with generating vector field
d(()))
(1(x)=1)) ox (x1) (x - pe (21)) + (1)))))
(1(1)) txı + (1 - t)æu, o((x1)= =o
(10)
2
Zt 11 20.
(11)
and the following flow matching loss:
fecused (h.g((a))p(aa (lo0)))l (x1 -0o)ll
(12)
A.1.2 RIEMANNIAN CONDITIONAL FLOW MATCHING
For side chain torsion angles, we need define a flow on the torus [-T, Fortunately, all components of the flow matching framework can be computed in simulation-free manner on this simple manifold. We use the explicit Riemannian conditional flow matching (RCFM) loss derived
by Chen & Lipman| 2023):
(0) (Br.g((1))p(a)))
(1,, l/ db. l12
(13)
The norm flg on the tangent space T.M point on manifold M induced by the Riemannian metric y which is the standard inner product (u, v) (u, v), u, E T,M in this particular case.
Choosing the geodesic path to define the conditional flow a, = vhe (2o))=)) we can compute inter- mediate points in closed-form as
exPx.. ((1 k(t)) logzo (x11)
(14)
using the exponential and logarithm maps exp, (u) w(x + u) and log, (3) atan2(sin æ), cos(y x)) w(a) = ([a+m) + T) mod 2T) wraps values within the range [-T, тT) and is applied element- wise.
xz, y M and u. T.M.
17

---

Published as a conference paper at ICLR 2025
Here, additionally scheduler k(t), that satisfies k(0) and x(1) = 0. to control the rate at which the geodesic distance d between 2o and æ1 decreases (Chen & Lipman 2023 :
d(æ;, = x(1)d((3o æi).
(15)
Thus, we obtain the loss function
(x = Es, (t.g(( (/ (20e, t) deel?
(16)
with
tt = kt()) logan (x1).
(17)
Here, the learned vector field va(2,, t) represents vectors on the tangent plane.
A.1.3 MARKOV BRIDGE MODEL
The molecular graph consists of discrete entities (node and edge types) and can therefore not be easily modeled in the flow matching framework. While discrete diffusion formulations (Austin et al. 2021 Vignac al. (022) can be used in principle, decided to employ the Markov bridge model [[gashov et al. 2023) instead which is conceptually more similar to the flow matching scheme used for the continuous variables as it does require a closed-form prior.
The Markov bridge model captures the stochastic dependency between two discrete-valued spaces X' and y. It defines Markov process between fixed start and end points 20 æ and ø y. respectively, through a sequence of N + random variables (Zt=i/N) ii for which
p(=r|=0, #0+ ... St, #1 = y) =
#1 = y)
(18)
with AL = 1/N. Additionally, since the process is pinned at its end point, we have
p(z1 = y 21 y)
(19)
Each transition is given by
p(Z 21, 21 y) Cat( Qizi)
(20)
where 2, E 10 is a one-hot representation of the current category and Qr is a transition matrix
parameterised
Q: =
- - (1- y.
(21)
Any intermediate state of the Markov chain can be probed in closed form:
p(zz 21) Cat( Q: Q.-a:
(22)
with
Q. Q. T.= 3,1 + (1 - --1y
(23)
In this work, we choose a linear schedule for B = L which implies B, = B./3. t)/(1 - At).
The neural network 9 approximates y so that we can sample from the Markov bridge without know- ing the true final state. It is trained by maximizing the following lower bound on the log- likelioold qo the end point y given the start point x
log go (2)p)) -T' (p) DRL
DKL (pl 21 zt, 1))lao = (OMBM(O). (24)
A.1.4 TRAINING LOSS
Our overall loss function is a weighted sum of the previously introduced loss terms:
L AcunndLaunnt + Xx Aa LMMM. som A LMMIM bond-
(25)
18

---

Published as a conference paper at ICLR 2025
A.2 PREDICTIVE UNCERTAINTY ESTIMATES FOR REGRESSION PROBLEMS
The basic idea behind our uncertainty estimation approach is enable the neural network to not only output a single predicted value ý but approximate the distribution within which the true value y is likely to be found. Assuming this distribution is Gaussian, it can be specified with two parameters, mean and variance. To this end, we view the task of approximating the target variable y E R as maximum likelihood estimation with a Gaussian model for the error e y and maximise the probability density of the data under this Gaussian uncertainty model with Σ 82
ply: ú. Σ) Ny::0.62)
(26)
exp( (y 1) g))
(27)
(2x) det Σ
1
(2x)d/2gu exp(
11y /1(²)
(28)
282
where neural network approximates both y (x)) and ā ãs (a:) based on the inputs æ.
Maximising the log-likelihood yields
d
arg max ply: D. Σ)} arg max -
dlog ā log 2
1y =
(29)
262
d
arg max
log 0²
272 byy 01
(30)
2
d
arg minl log 6² +
28 y - 0113]
(31)
2
which motivates the loss function
d ((0) log 02
28 y 8/1
(32)
2
Intuitively, if the neural network encounters rare training sample, it might be uncertain about its prediction and anticipate a large error |yy " To keep the overall loss small regardless, it can now assign larger variance 6² this "out-of-distribution" sample. This means, according to model, the true value might be found anywhere within a Gaussian distribution with a large spread. Moreover, the additive term & log 8² ensures the network does not achieve small loss values simply by excessively increasing à²
Regularization In Eq. 32 the standard regression loss is recovered if 1. We find empirically that we can stabilize traiinn through L2 regularization of the newly introduced parameter d² as follows:
L(0)
d log 2
1 2
2
10 10 11
(33)
2
Regression loss with uncertainty estimation
Regularization of the uncertainty estimate
This approach can be justified probabilistically by imposing a Gaussian prior with mean and variance 1/A on 6². First, for simpler notation we redefine p(y; ÿ. Σ) as p(y]a²) = p(y; ú, Σ) Next, using Bayes's theorem, we can convert the likelihood of an observed data point y given the parameter 2 into posterior distribution:
p(o2ly) & p(y))3)p(o2).
(34)
NNote that we use the term "out-of oout-of-disten distribution" interchangeably with "uncertain" in the context of our generative model. Higher uncertainty implies that the sample is more likely to contain patterns dissimilar to the main modes of the data distribution empirically supported by the data in Figure
19

---

Published as a conference paper at ICLR 2025
With prior p(0²) (1, then solve the maximum a posteriori (MAP) problem
o MAP
arg maxx|logp(0*
(35)
a²
arg max maxxlogpp00 (log )²) log p(o2)]
(36)
a²
d
26 y - 011 + (0² 1)2
arg min
log ā²
(37)
2
2(1/A)
arg min d log 0²
أ11 + X 2002 2102- 1|3].
20 y
(38)
2
Applying this idea to the flow matching setup, obtain the loss in Eq. 1
A.3 MOTIVATION FOR THE FINAL UNCERTAINTY SCORE
In the flow matching setting, we integrate an ordinary differential equation (ODE) along a time- dependent vector field. However, the true data generating vector field is not known and we only have access to the estimated neural vector field vo (2,,, t) which introduces epistemic uncertainty into the process. We can model the true vector field u, (x( as a sum of the estimated vector field and normally distributed error term 6; ~ ~NO of (æt, f)
u, = By (z,, t) + Ex.
(39)
The newly introduced stochasticity can be described with stochastic differential equation (SDE) rather than the original ODE
da: = vy vo( l)dt + ca l)d
(40)
where B, is the Wiener process.
As a result of integration, the final data point is
ac = 20 + / vo(xs, s)ds+ + f co(as, s)dB..
(41)
We could find approximate numerical solution using the Fuler-Maruyama method:
B++ = 3. + Ug (axe l) + as(a,, () Aiz,
(42)
where NO. I). However, here opt the learned vector field in the standard way and integrate in a deterministic manner. This means we only have access to op(: 8) and (o(:, 8) along deterministic trajectory. If we assume alternative SDE parameterized by these values, we can write
ã L'' vo(s)ds+ / co(S)dB,
(43)
and compute mean and variance of the final variable af
E [ xo + S' vo()sss,
(44)
Var[ø:] = /''' oz (s)ds.
(45)
This means obtain the most likely data point according to our leamned model by following the predicted mean as in conventional flow matching. However, this view point additionally provides us with an estimated variance of the sample. Note that Eq. 43 represents purely hypothetical scenario created to motivate our final uncertainty score. More specifically, it does not describe actual DrugFlow sampling trajectories, which are fully deterministic except for the prior. Therefore [[ and Var [ē,] do not represent the mean and standard deviation of sampling outputs.
Note that cannot derive this formal solution if the stochastic process appears on both sides as in Eq. 41 hecause the integrands will contain expectations and variances well.
20

---

Published as a conference paper at ICLR 2025
A.4 PREFERENCE ALIGNMENT
Our preference alignment scheme is based on Direct Preference Optimisation (DPO), a technique proposed by Rafailov al. (2023) as more stable alternative reinforcement learning from human feedback methods aimed to optimise a model with respect to a target metric (i.e. reward function). Instead of explicitly learning the reward, Rafailov al.) (2023) proposed to fine-tune pre-trained (reference) model & on a synthetic dataset D = [[x h collected from its own samples. This dataset consists of pairs of winning and losing samples x; and æt, respectively, which are labeled according problem-speciffe preferences (i.e. using human feedback some other oracle). To align the method with these preferences, the authors introduce new (aligned) model initialised with 0 and further optimised with new loss that accounts for the provided preferences (Rafailov et al. 2023).
Initially proposed for aligning large language models, DPO was further adapted to diffusion models by Wallace et al. (0024). For noisy versions x; and zt (for clarity, we omit indices i) of the winning and losing data points, the diffusion DPO loss computed using the true transition kernel p and approximated transition kernels qo and 9.p of the reference and aligned diffusion models as follows,
LpPO-Diflusion (4)
log (-
DKL (p(x |a tll l2)))
DKI
t) (rx"))
DKI (pla 0) (
+ DKa.(((at | 1) (1 lx))))
(46)
We propose to apply this framework to the Markov bridge models. Using Eq. 24 we can derive:
(DPO-MMM(P) = - log o 3(CNiB - CMBM (0) - C.IBM (x) + (MBM(())
(47)
Here, 4% (M) and (+) are Markov bridge loss terms of the aligned model for winning and losing samples, respectively, and (0) and (0) are the loss terms for the pre-trained refer- ence model f with fixed parameters.
We then apply the same scheme the coordinate flow matching loss and define the multi-domain preference alignment (MDPA) loss as a weighted sum of the different loss components. Further- more, we scale the weighting constant by the sampled time E [0, 1]. as proposed in Cheng et al. 2024)
EMDDA (4) = log o (BBt()cccmmee BBteeeond + Aatom A Anom + Apona Abond)
(48)
where
A coord
Looord (1p) (0) - (ooord((p)) + Cooord
A atom MBM. atom ( CMBM. atom (0) CMUBM utorm (4) + L L atom (0),
! Aond MBM. bond bond((p) CMBM. bama( ChUBM, bond (1p) + CLaSM, bond (0).
Next, we introduce an additional regularization term, which is the scaled original loss (Eq. 25 applied to both winning and losing samples. The overall loss function is defined as the weighted sum of EMDRA and the regularization term:
LMDPA (4p AMDDA (MmPx + Ct) + (yp)
(49)
Setting AMDRA A = corresponds to simple fine-tuning of the model on the winning samples, which we use as a baseline in our experiments.
A.5 MODEL ARCHITECTURE AND TRAINING
Input graph definition While the computational graph of the generated small molecule must necessarily be complete that bond types can be generated freely, we improve the computational
21

---

Published as a conference paper at ICLR 2025
A
Ligand edges
Piotcin edgen
GH-GNN layer
Ligand
Protain
Edge Features
Fextunce
Featurizr
Feutuncee
Protcin
L-L
F-P
Time embedding
Haterogeneous
(GH-GNN)
Mesuage
Message
Mansage
Masuage
Soalrr festures
Vector fretures
All features
Uncorsainty
Atom lngits
Coordirasa
Now
Angle foww
Sond lngits
Somtllu
Proteir
C
Layutlorm
Dripoot
Layehoorm
15
41
Outpur
Sepp coppoomnetioon
Figure 6: Architecture of our backbone neural network. (A) Different input types are featurized independently and processed with E(3) -qquivaair heterogeneous graph neural network based on Geometric Vector Perceptrons (GVP) (Jing et al., (220). (B) One layer of the geometric het- erogenous GNN (GH-GNN). Messages are computed based source and destination node features and the corresponding edge features using GVPs. They are aggregated separately for appropriate destination node types and passed the output block. (C) The output block of GH-GNN layers contains another equivariant GVP module.
efficiency by removing edges between pocket residues or between residues and ligand atoms based on predefined cutoff distance (10A). Nodes in this graph correspond either to ligand atom or a pocket residue. The coordinates of the residue nodes are defined by the position of their C. atoms. To retain the full atomic information while adopting this coarse-grained representation for protein pocket (one computational node per residue), we include difference vectors to each atom of the residue in addition the Ca coordinate and amino acid type node input features similar to Zhang et al. (2023a). A schematic representation of the input graph showing different types of nodes and edges provided in Figure 1
Featurization We consider the atom types {C, N, O, B, Br, C, NH, N+, O-} where +/- indicate charges and NH is nitrogen atom with explicit hydrogen. all other cases, hydrogens assumed to be implicit following normal valence assumptions. Furthermore, DRUGFLOW generates single, double, triple, aromatic, and "None" bond types. FLEXFLOW additionally outputs five torsion angles (x1. X2, X3. X4. ) for each residue. Since not all angles are present in every residue we mask predictions where appropriate.
ligand nodes we include node-level cycle counts up to size 5 following Vignac et al. (2022) and Igashov et al. (2023).
Self-conditioning Self-conditioning (Chen et al. 2022) is a sampling strategy in which the neural network takes its previous prediction additional input during iterative sampling. Like previous works Yim et al. 2023b Stärk et al. 2023) we observe significant performance improvements using this technique.
Neural network Since our computational graph contains two distinct groups of nodes, ligand and residue, and four different kinds of edges, Aggan-ttoligaam (LxLL, ligand uigand-to-poeee to-pocket (LsPP), pocket-to- ligand (PxLL) and tocket-to-pocker (Pt-P), we use heterogeneous graph neural network architecture as depicted in Figure It performs message passing operations using separate learnable message functions for each edge type, and separate update functions for each node type. All these functions
are implemented with geometric vector perceptron (GVP) layers (Jing et al. 2020 202 to ensure equivariance to global roto-translations.
Number of nodes To choose a number computational nodes during sampling, which represent upper bound on the final number of atoms, compute the categorical distribution p( [[N] M)
22

---

Published as a conference paper at ICLR 2025
(histogram) of molecule sizes N given the number of residues M in the target pocket based on the training set, sample from it, and add 2 extra nodes to account for the expected number of virtual nodes.
Hyperparameters
Important model hyperparamelers are summarized in Table 4
Table 4: Model hyperparameters.
Model
Parameter
DRUGFLOW
DRUGFLOW-OOD
FLEXFLOW
DRUGFLOW-PA
Num. weights (M)
12.1
12.1
12.6
12.1
Model size (MB)
48.408
48.411
50.557
48,408
Training epochs
600
600
700
250
Virtual nodes Nmax
10
10
Sampling steps
500
500
500
500
Uncertainty head Flexible side chains
No
Yes
No
No
No
No
Yes
No
OOD A
Scheduler &
Preference alignment
B
100
coord
1
Asom
0.5
Abond
0.5
Aw
A
0.2
23

---

Published as a conference paper at ICLR 2025
B EXTENDED RESULTS
B.1 EXTENDED DISTRIBUTION LEARNING METRICS
In this section, provide visual comparisons of distributions of various molecular characteristics from Tables 2
and 3 First, we apply PCA to the molecular embeddings computed by ChemNet,
the neural network used to calculate the Fréchet ChemNet Distance (FCD) (Preuer al. 2018). We compare distributions of the first two principal components in Figure 7 Next, we compare distribu- tions of discrete data types (atom and covalent bonds) Figure Finally. we provide violinplots for continuous distributions of various geometric and chemical properties, binding efficiency scores and normalised interaction counts in Figures 10] and II Note that we remove outliers (beyond the 1st and the 99th percentiles) for better visibility. We additionally report Jensen- Messnn-Shannon divergence for various geometric and chemical quantities, following the methodology chosen other works (Guan al.) 2023a). To do this, compute histograms of the scores splitting them in 100 bins of equal sizes on ranges defined by the minimum and maximum values of the corresponding quantities in the training data. The results provided in Table
Pocket2Mol
DirIsBDD
TargetDiff
DrugFlow
4-
2-
Figure 7:
Distributions of the first two principle components molecule embeddings computed by the FCD neural network. In each plot, we compare training data (blue areas) and samples generated by different methods (black solid lines).
Atoms
Covalent bonds
0.8
0.7
Train data
Train data
Pocket2Mol
Pocket2Mol
0.6
DiffSBDD
0.6
DiffSBDD
0.5
TargetDiff
TargetDiff
DrugFlow
DrugFlow
0.4
0.3
0.2
0.2
0.1
0.0
0.0
N
O
S
SINGLE
DOUBLE
TRIPLE
AROMATIC
Figure 8: Distributions atom types (4 most popular types) and covalent bond types in the training data and samples.
24

---

Published as a conference paper at ICLR 2025
C-C
C=C
1.7
16
15
14-
1.3
12-
Samples Trin
12
2oc DIRSe3D Targetbiff Drugflow
Pocket2Mol DIfSSDDO Target Brugflow
Pocket2Mai DINSBDD TergeDDiff Drughtow
C-c-c
C-C-C
C-C-O
150
160
140
140
9 130
120
120
100
110
80
60
Pocket2Mol
TargetDiff DrugFlow
Nocket2MMll
TargetDiff Drugflow
TurgetDiff DrugFicw
Figure 9: Distributions of bond distances and angles. Training data is visualized in blue and samples in red. Outliers falling beyond the 1st and the 99th percentiles are removed for better visibility.
Retorable berds
0.6
g
02
Tram data
NocedtPe D460000 TangeDD Brughke
Pookdttte 1ea0D Targethi Drufhin
PecusiMM D15600 Tapel
Figure 10: Distributions of molecular properties. Training data visualized in blue and samples in red. Outliers falling beyond the 1t and the 99th percentiles are removed for better visibility.
0.2-
Samplas
60-
6.15
0.10
0.75
-0.5
3.0
0.00
UM51UD
Tergetorf Drag fow
Defa00 Tepetii Brgfice
DM55300 Tarpttof Dragl kn
PactaciMe
Teretidf Drigknn
Grina efficieccy
estackig
Hyskophubs
6.122
02-
6.000
100
2rdec2PPl DI5500 "anpoit Cragiher
Tacec2M 2415220
Rapetie Dogftn
PrubodMt EM15t0 Taeeiter Draghton
Figure 11: Distributions binding efficiency scores and normalised numbers of interactions. Train- ing data is visualized in blue and samples in red. Outliers falling beyond the 1st and the 99th percentiles are removed for better visibility.
25

---

Published as a conference paper at ICLR 2025
DEDRRD Phoke
Tangett
-30
30-11
20111
<10
<010
20
<10
<10
uniiao
CwC
C-C-C
CCO
log servolus
30
Teyette
-10
Dhaghen
Draffio
PrcketMell
or1b000
Tangitt1
Drugfion
DV11800
Tingo
Driffin
DI15600
Targetot
Drughhe
whether there is statistically significant difference between samples across various scores and methods. Here, we compare distributions of bond lengths and angles. We perform pairwise comparisons of distributions for all possible pairs of methods and report the resulting p-values. As shown in the last row of each matrix, the distribution differences between DRUGFLOW and other methods are statistically significant almost everywhere.
Figure 12:
Mann-Whitme U test
QED
SA
Di 300000000 hasca
<300
Dragtine sayoth
Dragtte
PhowitMal
Ta
Drughion
NowdaMer
DrugPon
NoNdPM
D115000
TarpoOot
Drughh
Rotelable
Doight Tapgoou ONIERD Pooeett
--10
locs (p
40
<30
<32
<10
PecwiOMal
PeculaMe
Leriaoo
Coigfio
Figure 13: MaineWhitney U test to assess whether there is a statistically significant difference be- tween samples across various scores and methods. Here, we compare distributions of molecular properties and binding efficiency scores. We perform pairwise comparisons of distributions for all possible pairs of methods and report the resulting p- p-values. As shown the last row of each matrix, the distribution differences between DRUGFLOW and other methods are statistically significant al- most everywhere.
26

---

Published as a conference paper at ICLR 2025
stacking
Hydrophobic
Prder
10% vevvluu
-10
CII500
13:11
-30-1
-300
<10
10
<33
<10
<100
10
Dronnzt
Drugtion
-50
Figure 14: Mann- Hoin-Whitney U test to assess whether there is a statistically significant difference be- tween samples across various scores and methods. Here, we compare distributions of normalised numbers of protein-ligand interactions. We perform pairwise comparisons of distributions for all possible pairs of methods and report the resulting p- p-values. As shown in the last row of each matrix, the distribution differences between DRUGFLOW and other methods are statistically significant al- most everywhere.
Table 5: lensen-Shannon divergence between distributions of continuous molecular data (bond dis- tances and angles), drug-likendd (QED), synthetic accessibility (SA), lipophilicity (logP) and num- bers of rotatable bonds (RB). The best result is highlighted in bold, the second best is underlined.
Top-3 bond distances
Top-3 bond angles
Molecular properties
Method
C-C
C-N
C=C
C-C=C
C-C-C
C-C-O
QED
logP
Rotatable bonds
POCKET2MOL
0.357
0.289
0.400
0.312
0.197
0.276
0.247
0,230
0.330
0.400
DIFSBDD
0.354
0.329
0.315
0.343
0.162
0,479
0.203
0.115
TARGETDIFE
0.236
0.219
0.321
0.317
0.222
0.254
0.138
0.470
0.148
0.083
DRUGFLOW
0.223
0.218
0.242
0.164
0.154
0.177
0.089
0.170
0.156
0.064
B.2 ABSOLUTE METRICS
In Table 6 we provide additional metrics evaluating overall quality of samples in absolute values. For reference, we also provide training set numbers and remind that is unreasonable to expect the model trained solely with the likelihood objective to substantially surpass metric values from the training data.
Table Absolute values of various quality metrics. The best result is highlighted bold, the second best underlined.
Passed filters, %
Binding efficiency
Method
Size
PoseBusters
REOS
Clashes f
Gmina
Vina J
Uniqueness
Novelty 1
MOL
13.024
0.867
0.408
0.926
0.318
-0.443
0.892
0.996
DIFFSBDD
24.397
0.380
0.110
0.696
0.231
-0.293
1,000
1.000
TIFGETDIFF
24.150
0.510
0.174
0.884
0.244
-0.345
0.989
0.998
DRUGFLOW
20.827
0.731
0.245
0.897
0.261
-0.351
0.956
0.997
Training set
23.648
0.948
0.248
0.977
0.274
-0.379
27

---

Published as a conference paper at ICLR 2025
B.3 ABLATION STUDIES
Here, we study the effect of using virtual nodes, uncertainty estimation and flexible side chains on the metries we reported in the main text. We repeat the evaluation procedure and compare four models:
DRUGFLOW, our base model from Tables and not the uncertainty head;
3
This model uses virtual nodes but
DRUGFI.OW (no virt. nodes), an identical model without virtual nodes;
DRUGFLOW-OORT an identical model trained with additional uncertainty head to be able detect out-of-distribution (OOD) samples. This model uses virtual nodes;
FLEXFLOW, identical model that additionally operates protein side chains. This model uses virtual nodes but does not have the uncertainty head.
As shown in Tables DRUGFL.OW. Notably, it remarkably improves Wasserstein distance on bond angles, QED, and SA scores. The model without virtual nodes demonstrates consistently worse performance across all metrics except lipophilicity. FLEXFLOW performs worse metrics related to pocket interac- tions. This result expected due to the much higher complexity of the flexible design task. While DRUGFLOW learns the conditional distribution of molecules given fixed (ground-truth) pockets, FLEXFLOW learns the joint distribution of molecules and side chain conformations.
and PRRGGFLoow-O0 demonstrates competitive performance with
To further contextualize the performance of our models, we include simple, unconditional baseline. For each test set protein, we randomly selected and docked 100 molecules from the 2.4M compounds in the ChEMBL database (release 34). We also repeated this procedure for CrossDocked training set molecules to provide another simple baseline.
The results in Tables 7 and highlight the importance of training curation for generating molecules with desirable properties. Most molecular properties (such as QED, SA score, and logP) deviate more from the CrossDocked training than DrugFlow's generated molecules. The FCD is substantially higher as well. A notable exception are more fundamental molecular features like bond lengths and angles for which the ChEMBL baseline is competitive. This can be explained by the universal physical rules all molecules must obey. Furthermore, Table Sable@ddemonstrates the impor- tance of pocket-conditioning Because ChEMBL molecules have been selected without taking into account the structure of the binding pocket, the distributions of most interaction-related features are matched consistently worse by this baseline than any of the Dluw variants.
Lastly, we also study the performance of DRUGFLOW depending the number of training epochs. As shown in Tables I0 and IT, the performance varies with training and is generally improving (Validity, FCD, Rings, RB), can be expected. Some other properties such atom types, bond types and geometries (distances and angles) are however learned rather quickly, and more training alone does not
guarantee better results.
Table 7: Fréchet ChemNet Distance and Paisen-Shannon divergence between distributions of dis- crete molecular data. The best result is highlighted in bold, the second best is underlined.
Method
FCD
Atoms
Bonds
Rings
DRUGFLoW
4.278
0.043
0.060
0.391
DRUGFLOW-OOD
4.328
0.054
0.019
0.374
DRUGFLOW (no virt. nodes)
5,380
0,050
0.093
0.411
FLEXFLOW
6.001
0.086
0.068
0.376
CHEMBL
9.255
0.124
0.038
0.377
CROSSDOCKED
0.578
0.066
0.003
0.186
28

---

Published as a conference paper at ICLR 2025
Table 8: Wasserstein distance between distributions of continuous molecular data (bond distances and angles), drug- drug-likeness (QED), synthetic accessibility (SA), lipophilicity (logP) and numbers of rotatable bonds (RB). The best result is highlighted bold, the second best underlined.
Top-3 bond distances
Top-3 bond angles
Molecular propertics
Method
C-C
C-N
C=C
C-C=C
C-C-C
C-C-O
QED
SA
logP
RB
DRUGFLOW
0.017
0,016
0.016
0.952
2.269
1.941
0.014
0.317
0.665
0.144
DRUGFLOW-OOD
0.021
0.029
0.021
0.683
1.412
1.478
0.005
0.140
0.796
0.176
DRUGFLOOW (no virt. nodes)
0.027
0.028
0.018
1.149
2.681
1.964
0.021
0.512
0.658
0.214
FLEXFLOW
0.019
0.019
0.017
1.021
1.731
1.937
0.029
0.231
1.115
0.497
CHEMBL
0.024
0.022
0.020
0.536
0.517
2.966
0.028
0.322
1.656
1.456
CROSSDOCKED
0.001
0.001
0.000
0.013
0.085
0.060
0.002
0.011
0.028
0.041
Table 9: Wasserstein distance between distributions of binding efficiency scores and normalized numbers of different protein-ligand interactions. The best result is highlighted in bold, the second best underlined.
Binding efficiency
Protein-ligand interactions
Method
Gnina
H-bond
H-bond (acc.)
H-bond (don.)
x-stacking
Hydrophobic
DRUGFLOW
0.028
0.013
0.019
0.012
0.007
0.006
0.036
DRUGFLOW-OOD
0.044
0.019
0.027
0.015
0.011
0.003
0.043
DRUGFLOW (no virt, nodes)
0.054
0.020
0.030
0.018
0.011
0.007
0.038
FLEXFLOW
0.073
0.040
0.077
0.052
0.025
0.012
0.203
CHEMBL
0.176
0.078
0.083
0.055
0.028
0.003
0.061
CROKED
0.071
0.061
0.019
0.081
B.4 UNCERTAINTY ESTIMATION
8000
6000
Count
4000
2000
0.825 0. 0.850 0.875 0.900 0.925 0.950 0.975
Uncertainty scores
Figure 15: Distribution of the uncertainty scores on the test set.
29

---

Published as a conference paper at ICLR 2025
Table 10: Dependency of the evaluation results on the number of training epochs. Validity, Fréchet ChemNet Distance and Jensen-Shannon divergence between distributions of discrete molecular data. The best result is highlighted in bold, the second best is underlined.
Epochs
Validity
FCD
Atoms
Bonds
Rings
100
0.785
6.064
0.044
0.070
0.476
300
0.807
5.775
0.078
0.045
0.416
500
0.885
5.154
0.067
0,063
0,416
600
0.891
4.278
0.043
0.060
0.391
Table 11: Dependency of the evaluation results on the number of training epochs. For all scores, Wasserstein distance between the corresponding distributions reported. The best result high- lighted in bold, the second best is underlined. RB: number of rotatable bonds.
Top-3 bond distances
Top-3 bond angles
Molecular properties
Epochs
C-C
C-N
C=C
C-C-C
C-C-C
C-C-O
QED
SA
logP
RB
100
0.040
0.035
0.025
1.021
2.035
1,811
0.022
0.655
0.344
0.719
300
0.015
0.020
0.014
0.508
1.533
1.372
0.055
0.589
1.260
0.515
500
0.016
0.024
0.018
0.883
1.417
1.417
0.014
0.438
0.487
0.347
600
0.017
0.016
0.016
0.952
2.269
1.941
0.014
0.317
0.665
0.144
B.5 VIRTUAL NODES
shrting
37 startigg
47 sarting
Figure 16: Samples for the pocket IL3L with varied input size.
30

---

Published as a conference paper at ICLR 2025
Table 12: Dependency of the evaluation results on the number of training epochs. Wasserstein distance between distributions of binding efficiency scores and normalized numbers of different protein-ligand interactions. The best result is highlighted in bold, the second best underlined.
Binding efficiency
Protein-ligand interactions
Training epochs
Vina
Gnina
H-bond
H-bond (acc.)
H-bond (don.)
s-stacking
Hydrophobic
100
0.038
0.019
0.028
0.016
0.012
0.007
0.028
300
0.039
0.035
0.019
0.011
0.008
0.004
0.075
500
0.056
0.029
0.024
0.011
0.013
0.005
0.041
600
0.028
0.013
0.019
0.012
0.007
0.006
0.036
31

---

Published as a conference paper at ICLR 2025
B.6 DISTRIBUTION OF SIDE CHAIN ANGLES
ARG
ASN
ASP
X2 (rad)
2
0
-
-4
GLN
GLU
HIS
4
X2 (rad)
2
0
-2
ILE
LEU
LYS
X2 (rad)
2
MET
PHE
PRO
X2 (rad)
2
0
-2
- 0
X1 (rad)
TRP
TYR
4
X2 (rad)
2
0-
Samples Training set
-2
-4-
-4
-2 0
2
-4 -2 0 2
Xi (rad)
X1 (rad)
Figure 17: Distributions of X1 and X2 angles for the 14 amino acids that have at least two side chain torsion angles. We compare FLEXFLOW samples to the bound pocket conformations from the training set.
32

---

Published as a conference paper at ICLR 2025
B.7 ADDITIONAL PREFERENCE ALIGNMENT RESULTS
Traihing set
0.6
Reference
25-
Preference alggnmett
Preference alignment (comb.1
0.5
0.4
Density
0.3
2-
0.-
6.0
0.2
0.4
0.6
0.5 1.0
00.5
-0.6
-0.4 -0.2
0.0
QED
SA
Vina efficiency
Figure 18: Preference alignment shifts property distributions. Distributions of QED, SA, and Vina efficiency values with and without preference alignment. Gray shaded areas represent the training set distributions, and green areas show distributions from DRUGFLOW as the reference model. Solid red lines indicate the distributions for preference-aligned models for each specific property, while the dashed red line shows the distribution for the model aligned with combined preferences. Preference alignment leads to significant shifts toward more desirable values in QED, SA, and Vina efficiency across all metrics.
Table 13: Preference alignment comparison. Performance comparison of preference-aligned (PA) models vs. fine-tuned (FT) models on REOS, QED, SA, Vina efficiency, and combined preference pairs, using DRUGFLOW as the baseline. We also report Vina scores (after local minimization). even though our models were not directly optimized for these scores. Values for AliDiff (Gu et al. 2024) and DecompDPO (Cheng et al. 2024) are reported by their authors (SA values were mapped to the original scale using SA = 10 norm). The authors of AliDiff provide sampled molecules and Vina scores, which derlined values indicate the second- best. Molecular validity, uniqueness, novelty, and PoseBusters success rates also reported. The preference-aligned models achieve the highest performance the target metrics, with combined preference alignment models ranking second. Preference align- ment models show a 10-20% drop in molecular validity compared less than 5% for fine-tuned models. PoseBusters measures robustness against common failure modes of generative models (But- tenschoen 2024
for re-evaluation F Bold values indicate the best performance, and un-
Molecular properties
Interactions
Method
Valid.
Uniq.
Nov.
PoseB.
REOS
QED
SA L
Vina eff.
Vina min
Baseline
0.89
0.95
1.00
0.73
0.25
0.53
3.49
-0.35
-6.67
FT (REOS)
0.87
1.00
1.00
0.62
0.55
0.57
3.85
-0.29
-7.35
(QED)
0.85
1,00
3.88
-0.30
-7.33
FT (SA)
0,86
1.00
1.00
0.58
0.24
0.50
3.45
-0.30
-7.41
(Vina efficiency)
0.77
1,00
4.37
-8.25
FT (Comhined)
0.84
1.00
1,00
0.57
0.51
0.58
3.74
-0.30
-7.47
PA (REOS)
0.79
0.97
1,00
0.67
0.88
0.65
3.90
-0.37
-7.08
(QED)
PA (SA)
0.76
0.93
1,00
0.45
0.29
0.57
2.39
-0.35
-6.50
PA (Vina efficiency)
0.65
0.95
1,00
0.37
0.29
0.56
3.83
+0.45
-7.78
PA (Combined)
0.73
0.94
1,00
0.69
0.67
2.45
-0.40
-7.18
AliDil evaluation)
0.92
0.99
1.00
S 0.26
0.20
0.50
4.92
-0.34
-7.87
AliDiff (Vina)
0.50
4.87
-8.09
(Vina, SA)
AliDiff (Vina, QED)
0.51
4.87
-8.01
DecompDPO (Vina)
0.48
4.06
-8.49
DPC (Vina, QED)
4.24
-7.93
Training set
1.00
0.95
0.25
0.53
3.23
-0.38
-8.29
'For evaluating AliDiff, molecules sampled for the CrossDocked test set were retrieved from For PoseBusters, only bocket-independen checks were completed, pocket information unavailable in the provided samples. Vina efficiency scores were derived from the reported Vina minimization
33

---

Published as a conference paper at ICLR 2025
REOS
Vina efficiency
Training
Bavdiin
-00
Finetune- couth -ppth
100 10000--00 Nocooall.o
-30
Brainng Aar
Boweline
Nise sxong
Phat atgen
Pt (cends
Raining Aey
Alsseinar
Piin Kune
align
Al rcamds
Brening ser
Alwilne
tame
algen
PA Ecemb
Brahing Svet
Buseinn
Prie-tund
alige
ccomes
Figure 19: Mann-Whinney U test to assess whether there is statistically significant difference between different preference alignment methods for REOS, QED, SA, and Vina efficiency scores reported in Figure S We perform pairwise comparisons of distributions for all possible pairs of methods and report the resulting p-values. As shown the last two rows of each matrix, the distri- bution differences between the preference aligned model and other methods statistically signifi- cant everywhere.
B.8 PREDICTION ERROR
10
100
Fred
Figure 20: Distribution of the prediction error of trained DRUGFLOW model. The values computed as (z - (x) zo) for one training batch. (A) All vector components are treated as independent (ID) samples. (B) Joint distributions for pairs of error components.
34

---

Published as a conference paper at ICLR 2025
B.9 DISTRIBUTION OF THE JOINT QED, SA, LOGP AND VINA SCORES
Pocket2Wol.550
DIMBOCCJED,
TargetDitt _5 C242
DrugFlow: 150 0.099
339119
haangg ba
0.003
Den500
0.301
Drugfow
Tra sdta
399115
60045
1.4155
0.005
13371
0/003
0.001
13302
3030
18:0
2008
4600
1900 3000
1603
(hh 10-00 legh.clim))
07.30 dAxiighh
Figure 21: Histograms of the joint distributions of QED, SA, logp and Vina efficiency scores (10 bins per score). We show top-5000 bins the rest is not populated.
B.10 SIGNIFICANCE OF DIFFERENCES IN WASSERSTEIN DISTANCE AND JENSEN-SHANNON DIVERGENCE
To evaluate the variability and sta- tistical significance of sample-based Wasserstein distances and Jensen- Shannon divergences, we generated 20 bootstrap samples, each contain- ing 500 data points (5 samples per test target). These samples were used to compute distances to the training set, in Section 3.1 Tables 16 provide the mean values and standard deviations of these distances over the bootstrapped samples. Statistical sig- nificance of the DrugFlow's 's superior
Table 14: Sample mean and standard deviation the
Jensen-Shannon divergence between distributions of molec- ular data, as presented in Table [ The final column shows the Jensen-Shannon divergence for the joint distributions of four molecular properties: QED, SA, LogP, and Vina effi ciency. Standard deviations are provided in brackets.
Method
Bonds
Rings
JSD
PoCKEr2MoL
0.082 (0.004)
0.045 (0.006)
0.491 (0.01 I)
0.226 (0.011)
DIRSBDD
(0.008)
(0.005)
(0.007)
(0.005)
TARGETDIFF
0.078 (0.004)
0.240 (0.011)
0.649 (0.007)
0.247 (0.010)
DRUGFLOW
(0.005)
(0.005)
(0.006)
(0.006)
performance is verified with Student's 's t-test, as shown in Figures 22
24
Table 15: Sample mean and standard deviation of Wasserstein distances reported in Table
Top bond
Molecular properties
Mechod
C-C
C-N
C=C
C-C+C
C C-C-C
C.C.O
QED
SA
JogP
RB
2ock
0.050 (0.001)
0.004 (0.001)
0.446 (0.002)
2.174 (0.122)
2.967 (0.362)
3,057 (0.254)
0.073 (0.004)
0.576 (0.077
1.210 co.035)
2.861 (0.07))
DivSHIDD
0.411 (0/001)
0.009 (0000))
0.92 (0:412)
A 0 (91)
K 165 (0.250)
(0.101)
(108 (0.008)
1.591 (0.0300
0 (0.047)
8915 vo. 188)
TAKGETDuS
0.013 (0/00))
0.019 (0-001)
0,077 (0:077)
(0:355)
(0:380
0L053 (0.017)
(0.00)
0.000 101)
DRUeFLow
(0.001)
8.017 (0.001)
(0,002)
0.957 (0.083)
(0.335)
1,972 (0.292)
(0.005)
0.330
0.682 (0.056)
0.114 (0.053
Table 16: Sample mean and standard deviation of Wasserstein distances reported in Table
Binding efficiency
Protein-ligand interactions
Method
Vina
Gnina
H-boad
H-bond (acc.)
H-bond (don.)
x-stacking
Hydrophobic
POCKET2MOL
0.064 (0.005)
0.044 (0.005)
0.040 (0.003)
0.026 (0.002)
0.014 (0.001)
0.008 (0.001)
0.028 (0.008)
DIFFSBDD
0.086 (0.004)
0.043 (0.002)
0.047 (0.002)
0.030 (0.002)
0.017 (0.002)
0.011 (0.001)
0.044 (0.003)
TARGETDIFF
0.036 (0.007)
0.031 (0.004)
0.055 (0.012)
0.045 (0.009)
0.013 (0.005)
0.012 (0.00))
0.063 (0.024)
DRUGFLOw
0.028 (0.005)
0.013 (0.003)
0.019 (0.003)
0.012 (0.002)
0.008 (0.001)
0.006 (0.001)
0.037 (0.004)
35

---

Published as a conference paper at ICLR 2025
CAC
PodertMle
Detzeo
10:
Cragh
10
12:9
ursico
Terilti
Leught
Degget
Dough low
Pockactieot
DragTto
C-CWCC
C-C-O
PoderMMe
DeIzeo
10
10
13
Crogfe
14
PecoiiMal
ursico
Terelti
Drugh
PeciaiG/Mll
Targeto
Doughtow
DINAALO
DrugFlow
-30
Figure 22: Statistical analysis of Wasserstein distance differences using Student's 's t-test across vari- ous metrics and methods for bond lengths and angles. Wasserstein distances were computed using 20 bootstrap samples, each containing 500 datapoints. A two-sample t-test was conducted to determine whether the differences of the sample means of the Wasserstein distances were statisti- cally significant. The p-values reported log-scale and visualized heatmap, As illustrated the last row of each matrix, the differences between DRUGFLOW and other methods are statisti- cally significant in nearly all cases.
36

---

Published as a conference paper at ICLR 2025
QED
log
Pecket
Kinn
10
Dose SN
Draghes
priicn
Vegiet
Datintoo
Tangeni
Rocatable bonds
Grina
10 10
10 10
Draaf
Vegetti
Driget
Figure 23: Statistical analysis of Wasserstein distance differences using Student's t-test across molecular property distributions, including QED, SA, LogP, number of rotatable bonds, and Vina/Gnina efficiency. As in Figure 22 p-values are displayed log-scale heatmap, with the last row indicating that distance differences for DRUGFLOW are statistically significant all cases.
Hibond
estacking
hydroshobic
Todkc
arscco
10
Segedder
DrgFow
Seyett
Figure 24: Statistical analysis of Wasserstein distance differences using Student's t-test across nor- malised numbers of protein-ligand interactions. As in Figure 22 p-values are displayed on log- scale heatmap, with last row indicating that distance differences for DRUGFLOW are statistically significant in all cases.
37


Reviews:
[{"review_id": "0Cb9guYdJl", "reviewer": "Reviewer_6xpR", "summary": "The paper introduces DRUGFLOW, a generative model designed for structure-based drug design. It seamlessly combines continuous flow matching with discrete Markov bridges to capture the chemical, geometric, and physical characteristics of three-dimensional protein-ligand data. The model provides an uncertainty estimate to detect out-of-distribution samples and employs a joint preference alignment strategy to guide sampling towards desirable metric values. Furthermore, the paper extends the model to explore protein conformational landscapes by concurrently sampling side chain angles.", "questions": "1. FlexFlow samples side chain configurations additionally—does pocket data exist pre- and post-binding? If absent, how does this approach differ from treating the pocket as context?  \n2. In line #134, on what basis is the error assumed to be normally distributed?  \n3. During training, while virtual nodes are added to each sample, are virtual bonds similarly included?  \n4. How does the paper synthesize preference pairs, as noted in line #206?  \n5. Why does the author evaluate bonds using the Wasserstein distance in Table 1, whereas other studies [1] and [2] apply KL and Jensen-Shannon divergences?  \n6. Why does the reported Wasserstein distance for QED, SA, and logP differ from those in Table 1 of paper [3]?\n\n[1] Xingang Peng, Shitong Luo, Jiaqi Guan, Qi Xie, Jian Peng, Jianzhu Ma, \"Pocket2Mol: Efficient Molecular Sampling Based on 3D Protein Pockets\"\n[2] Jiaqi Guan, Wesley Wei Qian, Xingang Peng, Yufeng Su, Jian Peng, Jianzhu Ma. \"3D Equivariant Diffusion for Target-aware Molecule Generation and Affinity Prediction\"\n[3] Arne Schneuing1, Charles Harris, Yuanqi Du, Kieran Didi, Arian Jamasb, Ilia Igashov, Weitao Du, Carla Gomes, Tom L. Blundell, Pietro Lio, Max Welling, Michael Bronstein. \"Structure-based Drug Design with Equivariant Diffusion Models\"", "limitations": "", "rating": 6, "confidence": 5, "soundness": 1, "presentation": 3, "contribution": 2, "strengths": "- The paper is articulated clearly and concisely.  \n- Figures and tables effectively present complex data and comparisons, enhancing accessibility for readers.  \n- The methodology is detailed thoroughly, supporting reproducibility.", "weaknesses": "The primary concern lies in the paper's technical soundness:  \n\n1. The treatment of the pocket is inadequately detailed—it is unclear whether the pocket is generated jointly with the molecule or used as context.  \n2. In Section 2.1, the uncertainty estimation involves several ambiguities:  \n   - In line #133, the assumption of the error being normally distributed is neither evident nor justified.  \n   - In line #143, $\\dot{x}_t$ is inaccurately referred to as a ground truth vector field; it should be considered a conditional vector field, given $x_0$ is known.  \n   - In line #987, $\\underset{\\theta}{\\max} $ is mistakenly used instead of argmax; also, maximizing Equation 30 is not equivalent to minimizing Equation 31, despite sharing the same minima—the loss surfaces differ.  \n   - The described technique resembles regularization, yet its role in quantifying an atom's uncertainty score remains unclear.  \n3. In Section 2.2, the concept of a virtual node demands more explanation; specifically, it's unclear whether virtual bonds exist when virtual nodes are incorporated.", "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": "MhyJylSVNn", "reviewer": "Reviewer_o8Zm", "summary": "This paper proposes DrugFlow, a multi-modal flow matching model for structure-based drug design. DrugFlow jointly models the distribution of ligand structures and receptor sidechain structures. It also includes an uncertainty estimate module and an adaptive ligand size selection module that address the issues overlooked by previous work. In addition, preference alignment techinique is used for property optimization, which increases the value of this work. Overall, this is a nice work that orchestrated various machine learning techniques, which are all well justified in the context.", "questions": "See Weaknesses", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The model considers side-chain flexibility, which is critical in ligand docking and design as receptors are mostly non-rigid. The side-chain flexibility issue has also been overlooked in previous SBDD methods until this work, to the best of my knowledge.\n- This model provides an estimate of uncertainty, which is improtant in molecular modeling area and can increase the practicality of the method. Uncertainty estimation has been a common practice in structure prediction settings, but it has also been overlooked in the previous structure-based drug generation methods.\n- This work demonstrated the use of preference alignment to control the properties of generated molecules, which increases the value of the model as in SBDD, there are many properties other than receptor structures need to be considered.\n- All the techniques introduced in this work are well justified by clearly organized experiments (Section 3.2-3.5). Notably, uncertainties visualized in Figure 2 are very informative and I find unrealistic structures (long carbon chains with bifurcation) were assigned high uncertainties, which agrees with the intuition that such unrealistic structures are uncommon in the dataset.", "weaknesses": "- Does the evaluation presented in Section 3.1 consider side-chain flexibility? It seems that the DrugFlow and FlexFlow are separate variants and only the FlexFlow considers side-chain flexiblity.\n- If Section 3.1 does not model side-chain flexibility, why not? Did the authors consider jointly sampling both ligand structures and side-chain torsional angles?", "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": "76apeKzEr4", "reviewer": "Reviewer_67Pj", "summary": "This paper describes DrugFlow, a diffusion model for drug-like molecules in 3D. Contributions over existing work are 1) uncertainty estimates from diffusion model 2) an adaptive size selection method 3) protein conformation sampling module 4) a preference alignment optimization scheme.", "questions": "See weaknesses above. The common thread between the weaknesses is answering the question \"how well does the model work\". Anything the authors can do to answer this question would be helpful.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "This paper had a lot of strong positives but also some strong negatives. Starting with the positives:\n\n- Good knowledge of the field: unlike many ML papers in this area, this work has no statements about drug discovery that seemed to portray an embarrassing lack of domain knowledge on behalf of the authors. I also agree with the assessment that many works train models for distribution matching and then evaluate them for optimization, which does not make sense\n- The end-to-end uncertainty estimate is a really nice idea (even though it isn't clear that it works well, see below)\n- The ability to add or delete models during generation is a nice idea and seems to work reasonably well\n- I liked that the authors tested a wide range of tasks in the experiments section\n- Presentation of the paper is really good, definitely in the top 10%. Regardless of the other criticisms I raise below, I can tell that the authors crafted the manuscript very well", "weaknesses": "In my opinion, the biggest weaknesses of this paper all come from the experiments. I've organized them under the following headings\n\n### You might not be measuring the right things\n\nEssentially all metrics in the paper are about how well the distribution of molecules generated by the model matches the training distribution. However:\n\n- Only _marginal_ (1D) distributions seem to be measured, rather than _joint_ distributions of properties (i.e. does the joint distribution of SAscore and logP look the same between training and test). In general, matching the marginal distribution _does not_ imply that the joint distribution matches. Figure 4 is an exception to this.\n- QED/logP/SA are all very simple quantities which _do not depend on the 3D structure in any way_. The significance of matching these values is not very clear to me.\n\n### Complete disregard for statistical variation => significance of results is unclear\n\nAlmost every quantity estimated is estimated from a finite sample of generated molecules (including Wasserstein distances, JS-distances, coverage of chemical space). This means that all quantities in tables are _statistical estimates_ with finite-sample variation. Moreover, there is additional variation due to the randomness of model training, etc. This variation is not accounted for in any of the Tables (as far as I can tell), making the claims of performance differences poorly supported. I think the paper needs to include measures of variation and/or statistical significance tests to qualify its claims.\n\n### Significance aside, performance over the baselines is unclear\n\nFirst, the performance of all baselines presumably varies with training. Presumably, more training would give a closer match. Did the authors re-train the baseline models themselves or use a pre-trained checkpoint? Were all models trained a similar amount? Even if the results are statistically significant for a given pair of models, I think it would help to know how much these differences change with training. Perhaps include a plot or table showing how performance changes with training size? (not a specific request, feel free to provide a similar but different piece of evidence if you think it is more appropriate)\n\nSecond, it would be helpful to include more baseline models, particularly ones using 2D (ie graph-based) methods. A simple 2D method could be a random perturbation of the SELFIES string for a molecule in the training set. Another option could be an RNN trained on SMILES strings. For any metrics that require 3D coordinates, they could be generated using force fields from rdkit (or some comparable method). I would not expect these baselines to be state-of-the-art, but knowing the performance of simpler methods helps contextualize the performance of more complex ones.\n\nThirdly, it seems that no baselines were run for the preference optimization experiment?\n\nFinally, it is unclear that the OOD performance described in section 3.2 is _practically_ significant. The uncertainty estimates seem to vary by only a tiny amount (I see that the color bar in Figure 2B has a range of only 0.008). This seems to suggest that the model is not actually very well-calibrated?", "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 (Poster)