RapidDock: Unlocking Proteome-scale Molecular Docking

## Abstract

Abstract Accelerating molecular docking – the process of predicting how molecules bind to protein targets – could boost small-molecule drug discovery and revolutionize medicine. Unfortunately, current molecular docking tools are too slow to screen potential drugs against all relevant proteins, which often results in missed drug candidates or unexpected side effects occurring in clinical trials.
To address this gap, we introduce RapidDock , an efficient transformer-based model for blind molecular docking. RapidDock achieves at least a 100 × 100\times speed advantage over existing methods without compromising accuracy.
On the Posebusters and DockGen benchmarks, our method achieves 52.1 % percent 52.1 52.1\% and 44.0 % percent 44.0 44.0\% success rates ( RMSD < RMSD absent \text{RMSD}< 2 Å angstrom \mathrm{\SIUnitSymbolAngstrom} ), respectively.
The average inference time is 0.04 0.04 0.04 seconds on a single GPU, highlighting RapidDock ’s potential for large-scale docking studies.
We examine the key features of RapidDock that enable leveraging the transformer architecture for molecular docking, including the use of relative distance embeddings of 3 3 3 D structures in attention matrices, pre-training on protein folding, and a custom loss function invariant to molecular symmetries.

## 1 Introduction

Accelerating the drug discovery process will revolutionize medicine. As most novel drugs are small molecules (Kinch et al., ), various deep learning methods have
been proposed to streamline the process of docking such molecules to druggable protein targets (Abramson et al., ; Corso et al., ; ; Qiao et al., ).
While impressive, none of these methods is both *accurate* and *fast*.

In fact, to obtain a comprehensive view of its effects, a molecule needs to be screened against thousands of proteins (Sjöstedt et al., ).
However, state-of-the-art methods
(Corso et al., ; ; Abramson et al., ; Qiao et al., )
report run-times
on the scale of seconds per protein on a single GPU.
Consequently, screening a relatively small database of one million molecules (Yu et al., ) against the *proteome*, i.e.,
all proteins in the human body, would take years, even with hundreds of GPUs.
Such time-frames are unacceptable in the drug development process.

To address this challenge, we introduce *RapidDock*, a transformer-based model that performs molecular docking in a single forward pass, in hundredths of a second on a single GPU.
RapidDock performs blind docking, using unbound, possibly computationally folded proteins, so it can be applied to unexplored protein targets.
Given the $3$D structure of a protein and a molecule, our method predicts all pairwise distances in the resulting protein-molecule complex, including the molecule’s atom-atom distances and atom-amino acid distances.

RapidDock achieves success rates, i.e., $\text{RMSD}<2$\mathrm{\SIUnitSymbolAngstrom}$$, of $52.1\%$ on the Posebusters benchmark (Buttenschoen et al., ) and $44.0\%$ on the DockGen benchmark
(Corso et al., ),
with 0.03 and 0.05 seconds average runtimes, respectively.
Our method significantly outperforms the recent
DiffDock-L (Corso et al., ) (22.6% success rate, 35.4 seconds on DockGen) or the
NeuralPLexer (Qiao et al., ) (22.6% success rate, 3.77 seconds on Posebusters).
Because of its accuracy and speed, RapidDock can enable novel use cases and research directions.
For example, with RapidDock, docking ten million molecules to all human proteins on a cluster with $512$ GPUs would take nine days, compared to about 20 years with DiffDock-L or even 200 years required with a computationally-intensive method like AlphaFold-3 (Abramson et al., ).

We supplement the analysis of RapidDock’s performance with a thorough study of design choices, which include: operating on relative distances,
pre-training on a protein folding task, and a special treatment of multiple molecular conformations in the loss function.

In summary, our contributions are as follows:

- •
We introduce RapidDock, a first-in-class transformer-based model capable of accurate high-throughput molecular docking.
- •
We run RapidDock on the most challenging benchmarks.
The model docks $52.1\%$ of examples with $\text{RMSD}<2$\mathrm{\SIUnitSymbolAngstrom}$$ on Posebusters, and $44.0\%$ on DockGen, which places it among the most accurate methods.
Importantly, inference on a single GPU takes only $0.04$ seconds on average, making RapidDock at least $100\times$ faster than comparable methods.
- •
We provide ablations and describe the design choices that led to RapidDock’s success.

Figure: Figure 1: RapidDock architecture overview. The molecule is represented by a sequence of its atoms and the corresponding matrix of distances. The protein is represented by its amino acid sequence and its matrix of distances. Learnable embeddings of these distance matrices are added to the attention matrices. Additionally, the model uses ESM-2 embeddings (Lin et al., ) to improve its protein representation and embeddings of atom charges to improve its molecule representation.
Refer to caption: /html/2411.00004/assets/x1.png

## 2 Method

RapidDock, illustrated in [Figure 1](#S1.F1), is based on the transformer encoder architecture (Vaswani, ) and returns the positions of atoms in the docked molecule.
RapidDock represents the protein and molecule spatial structures by distance matrices whose embeddings are added to the attention matrices. Because RapidDock operates on distances, it is *equivariant*, i.e. invariant to translations or rotations, which we found crucial for accurate docking. Below we present the details.

### 2.1 Molecule and protein representation

The input to RapidDock consists of a sequence of tokens corresponding to molecule atoms and protein amino acids. The 3D structures of the protein and molecule are represented by learnable embeddings of their respective distance matrices. Additionally, we include the information on the molecule’s atom charges. We also utilize amino acid embeddings obtained from the protein language model ESM-2 (Lin et al., ).

##### Atom and amino acid dictionary

We have $26$ amino acid tokens ($20$ standard amino acids and $6$ tokens for non-standard or unknown residues) and $56$ atom tokens (all atoms likely to be found in drug molecules, i.e., elements lighter than uranium, excluding noble gases or reactive alkali metals).

##### Ligand atom charges

To account for electrostatic interactions, we include information about atom charges in the molecules. For each atom we include its partial charge (Gasteiger charges), which is the distribution of the molecule electron density spread over the ligand atoms. The values of charges are embedded (see below) and added to the atom embeddings.

##### Protein language embeddings

To improve the protein representation, we use embeddings from the 650M ESM-2 model Lin et al. (), one embedding per amino acid. These embeddings are added to the amino acid embeddings (to match the dimensionality, we use a linear projection layer).

##### Molecule rigid distance matrix

The ligand spatial structure is represented by the distance matrix ${\bm{D}}^{l}$ between its heavy atom coordinates ${\bm{x}}^{l}\in\mathbb{R}^{L\times 3}$, where $L$ is the number of heavy atoms in the ligand. Only the fixed distances across the molecule’s possible conformations are recorded, and others are denoted by a special value of $-1$. This distinguishes the rigid and moving parts of the molecule, which is important for modeling its allowable movements. To obtain the matrix, we generate 96 molecule conformations in total, using two methods: 64 from RDKit’s built-in algorithm based on MMFF optimization and 32 from the torsion-based algorithm described in (Zhou et al., ). The distances are averaged across all conformations and only elements with a standard deviation below 0.3 $\mathrm{\SIUnitSymbolAngstrom}$ are recorded.

##### Protein distance matrix

To represent the protein’s 3D structure, we extract the positions of alpha-carbon atom for each amino acid ${\bm{x}}^{p}\in\mathbb{R}^{P\times 3}$, where $P$ is the number of amino acids in the protein, and calculate the distance matrix ${\bm{D}}^{p}=\left(\|{\bm{x}}^{p}_{i}-{\bm{x}}^{p}_{j}\|\right)_{1\leq i,j\leq P.}$

### 2.2 Model architecture

RapidDock model is based on standard transformer (in our implementation, we use Mistral (Jiang et al., ) architecture), with full attention mask used instead of the causal one due to the non-autoregressive nature of the docking task. Below, we describe the modifications to the model: learnable attention biases based on the distance matrices, attention scalers, and learnable charge embeddings.

##### Distance biases matrices

Given the input distance matrices ${\bm{D}}^{l}$ for ligand and ${\bm{D}}^{p}$ for protein, we construct a block matrix ${\bm{D}}=\begin{bmatrix}{\bm{D}}^{l}&-1\\
-1&{\bm{D}}^{p}\end{bmatrix}$. Further, following (Raffel et al., ), we discretize the distances into $257$ buckets with:

$$ $b(x):=\begin{cases}\left\lfloor 8x\right\rfloor+1&\text{if }0\leq x\leq 16,\\ \min\left(256,\left\lfloor 128+128\frac{\ln\left(\frac{8x}{128}\right)}{\ln\left(6\right)}\right\rfloor\right)+1&\text{if }x>16,\\ 257&\text{if }x=-1.\end{cases}$ $$

Intuitively, the discretization is finer for small distances and coarser for larger ones, for which high accuracy is not necessary. The discretized values are used to construct distance bias matrices. For each attention head $l$ we put

$$ ${\bm{B}}^{l}=\left({\bm{E}}_{b\left({\bm{D}}_{i,j}\right),l}\right)_{1\leq i,j\leq L+P,}$ $$

where ${\bm{E}}\in\mathbb{R}^{257\times H}$ is a learnable embedding matrix and $H$ is the number of attention heads.

##### RapidDock attention

We modify the attention mechanism in two ways (marked in blue in ([1](#S2.E1))). First, we multiply the attention scores corresponding to input pairs with known distances (i.e. ligand-ligand within a rigid part and protein-protein) by a learnable scalar $s_{m}$, one for each layer $m$.
Second, we inject distance information by adding $z_{m}\cdot{\bm{B}}^{l}$, where $z_{m}$ is another learnable scalar for each layer $m$ and ${\bm{B}}^{l}$ is the distance bias matrix. The final attention formula for the $l$-th head in the $m$-th layer is defined as follows:

$$ $\text{Attention}({\bm{Q}},{\bm{K}},{\bm{V}})=\text{softmax}\left(\frac{{\bm{Q}}{\bm{K}}^{T}}{\sqrt{d_{k}}}{\color[rgb]{0,0,1}\odot{\bm{S}}_{m}+{\bm{z}}_{m}\cdot{\bm{B}}^{l}}\right){\bm{V}},$ (1) $$

where ${\bm{Q}},{\bm{K}},{\bm{V}}$ are the query, key, and value matrices, $d_{k}$ is the dimension of the key vectors and

$$ ${\bm{S}}_{m}=\left(\begin{cases}1&\text{if }{\bm{D}}_{i,j}=-1\text{ (unknown distance) }\\ {\bm{s}}_{m}&\text{otherwise }\\ \end{cases}\right)_{1\leq i,j\leq L+P.}$ $$

##### Charge embeddings

We discretize charge values into $22$ buckets using the following formula

$$ $c(x)=\left\lfloor 10\cdot\max(\min(x,1),-1)+11.5\right\rfloor,$ $$

with the 0 0-th bucket reserved for unknown charges. Our model learns an embedding matrix ${\bm{C}}\in\mathbb{R}^{22\times D}$, where $D$ is the dimension of the model’s hidden states.

##### Molecule and protein input

The model input is the concatenation: ${\bm{X}}=\left[{\bm{X}}^{l};{\bm{X}}^{p}\right]\in\mathbb{R}^{(L+P)\times D}$. For the $i$-th molecule atom, ${\bm{X}}^{l}_{i}$ is the sum of the learnable atom embedding and the charge embedding. For the $j$-th amino acid, ${\bm{X}}_{j}^{p}$ is the sum of the amino acid embedding and the ESM-2 embedding.

### 2.3 Training loss function

The model outputs predicted distance matrices, $\hat{{\bm{D}}}^{l}$ for the ligand, $\hat{{\bm{D}}}^{p}$ for the protein and $\hat{{\bm{D}}}^{lp}$ for ligand-protein distances. More precisely, a linear layer is applied to the transformer outputs, resulting in protein and ligand representations $\hat{{\bm{x}}}^{l}\in\mathbb{R}^{L\times H}$ and $\hat{{\bm{x}}}^{p}\in\mathbb{R}^{P\times H},$ where $H$ is the output dimension, equal to $16$ in our experiments.
The distance matrices are then defined as $\hat{{\bm{D}}}^{l}=\left(\|\hat{{\bm{x}}}^{l}_{i}-\hat{{\bm{x}}}^{l}_{j}\|\right)_{1\leq i,j\leq L}$, $\hat{{\bm{D}}}^{p}=\left(\\
|\hat{{\bm{x}}}^{p}_{i}-\hat{{\bm{x}}}^{p}_{j}\|\right)_{1\leq i,j\leq P}$, $\hat{{\bm{D}}}^{lp}=\left(\|\hat{{\bm{x}}}^{l}_{i}-\hat{{\bm{x}}}^{p}_{j}\|\right)_{1\leq i\leq L,1\leq j\leq P}.$ We apply the L1 loss to them during training. In order to keep the model focused on close interactions, we ignore the loss on terms that have both the predicted and ground-truth distances above $20$$\mathrm{\SIUnitSymbolAngstrom}$.
The total loss consists of three parts: ligand loss $\mathcal{L}^{l}$, protein loss $\mathcal{L}^{p}$ and docking loss $\mathcal{L}^{d}$:

$$ $\mathcal{L}=\mathcal{L}^{l}+\mathcal{L}^{p}+\mathcal{L}^{d}\text{, where }$ $$

$$ $\mathcal{L}^{l}=\sum_{i<j}|\hat{{\bm{D}}}^{l}_{i,j}-{\bm{D}}^{l}_{i,j}|\cdot\mathds{1}(\min(\hat{{\bm{D}}}^{l}_{i,j},{\bm{D}}^{l}_{i,j})<20),$ $$

$$ $\mathcal{L}^{p}=\sum_{i<j}|\hat{{\bm{D}}}^{p}_{i,j}-{\bm{D}}^{p}_{i,j}|\cdot\mathds{1}(\min(\hat{{\bm{D}}}^{p}_{i,j},{\bm{D}}^{p}_{i,j})<20),$ $$

$$ $\mathcal{L}^{d}=\sum_{i,j}|\hat{{\bm{D}}}^{lp}_{i,j}-{\bm{D}}^{lp}_{i,j}|\cdot\mathds{1}(\min(\hat{{\bm{D}}}^{lp}_{i,j},{\bm{D}}^{lp}_{i,j})<20),$ $$

where ${\bm{D}}^{lp}=\left(\|{\bm{x}}^{l}_{i}-{\bm{x}}^{p}_{j}\|\right)_{1\leq i\leq L,1\leq j\leq P}$ is the distance matrix between true atom and protein coordinates, whereas ${\bm{D}}^{l}$ and ${\bm{D}}^{p}$ are as defined above.

In case of molecules with permutation symmetries, such as benzene, a single docked pose might be represented by several sets of atom positions, and thus different distance matrices. In such a case, the loss function may not assign the smallest value to the correct conformation. To alleviate this issue, we apply the loss function to all of the possible permutations of the molecule and use only the best match (the one with the smallest loss value) in the backward pass. This prevents the model from having to guess the specific atom order used in the docked pose. This procedure is similar to the Permutation Loss introduced in (Gao et al., ), with our approach applying permutations to the ground truth (labels) instead of the predictions.

### 2.4 Ligand reconstruction

The last transformer layer outputs the predicted distance matrices. To reconstruct the 3-dimensional coordinates of the ligand atoms $\tilde{{\bm{x}}}^{l}\in\mathbb{R}^{l\times 3},$ RapidDock uses L-BFGS (Liu & Nocedal, ) algorithm with the following objective function defined on the predicted ligand-ligand and ligand-protein distance matrices:

$$ $\mathcal{L}^{r}=\sum_{i,j}|\|\tilde{{\bm{x}}}_{i}-{\bm{x}}^{p}_{j}\|-\hat{{\bm{D}}}^{lp}_{i,j}|\cdot\mathds{1}(\hat{{\bm{D}}}^{lp}_{i,j}<20)+\sum_{i<j}|\|\tilde{{\bm{x}}}_{i}-\tilde{{\bm{x}}}_{j}\|-\hat{{\bm{D}}}^{l}_{i,j}|\cdot\mathds{1}(\hat{{\bm{D}}}^{l}_{i,j}<20),$ $$

with initial guess for molecule atoms defined as the weighted mean of the four closest amino acids based on the predicted distance matrix $\hat{{\bm{D}}}^{lp}$. The computational cost of this operation is comparable to the transformer forward pass and is included in our runtime results. For details, see [Section A.5](#A1.SS5).

## 3 Experiments setup

### 3.1 Datasets

##### Training and test datasets

Following (Corso et al., ),
RapidDock is trained on PDBBind (Liu et al., ) and BindingMOAD (Hu et al., ) and tested on Posebusters (Buttenschoen et al., ) and DockGen (Corso et al., ) datasets. In addition, we augment the training dataset by computationally generated apostructures – unbound protein structures in the absence of a ligand – for about 30% of the training examples using AlphaFold-2 server, so that the model can see both the unbound and ground-truth bound proteins (also called *holostructures*).
The apostructures are only used for defining the embeddings of the input distance matrices, while the loss is always computed based on the holostructures.
This way, we allow the model to learn to deform the protein during the docking process.

##### Ligands preparation

The $3$D coordinates of ligands are read from their PDB files and their graph structure is matched with their CCD reference.
Entries for which such matching fails are filtered out.
Ligands with more than $128$ heavy atoms are filtered out.
Whenever calculating the ligand rigid distance matrix is impossible because of a failure in calculating conformations or force fields, such entry is removed.

### 3.2 Training Procedure

RapidDock’s transformer architecture has $16$ layers,
each with $4$ attention heads and a hidden size of $512$, resulting in approximately $60$ million parameters. We found that such "deep and narrow" architecture performed better than shallower-but-wider models with the same number of parameters.

Since similar physico-chemical principles govern both protein folding and ligand binding, we pre-trained the model on a protein folding task with the hope that it would improve the docking performance.
We used approximately $440\text{k}$ structures generated by AlphaFold-2 on the SwissProt protein database (Varadi et al., ; Gasteiger et al., ). More precisely, the model was trained to predict distances between amino acids, with a masking factor of 97% applied to the input distances, which effectively simulates the protein folding task.

Following pre-training, the model was fine-tuned on the molecular docking task on our training dataset.
In total, the training process took about 3 days on a machine with eight A100 GPUs. The details are described in [A.4](#A1.SS4)

## 4 Experiments

We evaluate RapidDock on two recent challenging benchmarks that were not part of the training dataset:
the DockGen benchmark (Corso et al., ), and
the Posebusters dataset (Buttenschoen et al., ). These benchmarks cover a wide range of structures across various protein domains.
We calculate the heavy-atom Root Mean Square Deviation (RMSD) between the predicted ligand and the crystal (ground-truth) ligand atoms. We report both the median RMSD and the percent of examples with RMSD $<$ $2$Å (denoted % RMSD $<$ $2$Å), the latter being a widely-used threshold for considering a molecule to be correctly docked (Bell & Zhang, ). A more detailed definition of RMSD is described in [A.2](#A1.SS2).

### 4.1 Main results

##### Accuracy and speed

We compare RapidDock, both in terms of accuracy and runtimes, to three recent state-of-the-art blind docking methods: AlphaFold-3 (Abramson et al., ), NeuralPLexer (Qiao et al., ), and DiffDock-L (Corso et al., ), see [Table 1](#S4.T1).
RapidDock is more accurate than DiffDock-L and about three orders of magnitude faster.
While the speed advantage over NeuralPLexer is less pronounced (about $100\times$ faster), RapidDock has an even larger accuracy advantage over that model.
Finally, RapidDock is a staggering four orders of magnitude faster than AlphaFold-3.
Although AlphaFold-3 achieves the highest accuracy, it is
computationally intensive and therefore less practical for high-throughput screening studies with thousands of compounds.

Overall, the experiments show outstanding performance of RapidDock, which offers excellent accuracy at unmatched speeds. The backbone transformer architecture is efficient and allows for performing the docking in a single pass. The model is also well-suited for applications to downstream tasks such as predicting the binding strength, e.g., by direct fine-tuning.

Table: Table 1: Comparison of state-of-the-art molecular docking methods. Details of how we obtained the results can be found in [A.3](#A1.SS3)

##### Unlocking Proteome-wide docking

To validate RapidDock’s potential for proteome-wide docking studies, we used our model to run inference for twelve well-studied drugs and toxins, and entire human proteome consisting of 19222 proteins from The Human Protein Atlas ([https://proteinatlas.org/](https://proteinatlas.org/)) (Uhlén et al., ). The whole process took on average 74 seconds per ligand on a machine with eight A100 GPUs.

### 4.2 Results on Apostructures

The Posebusters and DockGen sets contain bound protein structures (i.e., the holostructures), which is a potential data leak, because the models might hypothetically use the knowledge of the deformation of the binding pocket to improve their predictions.
Therefore, to better evaluate RapidDock’s practical applicability, we tested its performance on a challenging real-world scenario where the model performs inference given only computationally folded apostructures.
We generated such apostructures using AlphaFold-2 for all proteins in the test sets.
Since our model
was trained to predict distances based on both the holostructures and the apostrucutres, we hoped it would be relatively unaffected by any potential information leakage present in holostructures.

We used the apostructures to obtain the distance matrices embeddings, so there is no information leakage for predicting the distances in the resulting complex. To check whether these predicted distances are still accurate, we perform evaluation in the same way as before, using the holostructures only for reconstructing the molecule’s position. We found that this leads to only a slight deterioration of performance metrics..
The percentage of examples with RMSD $<$ $2$Å is $51.7\%$ on the Posebusters set, and $42.3\%$ on the DockGen test set, a decrease of only $0.4\%$ and $1.7\%$, respectively, compared to using holostructures in the input.

### 4.3 Ablation Studies

We performed ablation studies to assess the relative importance of RapidDock’s key features as shown in [Table 2](#S4.T2). We found protein language model representations (ESM-2 embeddings) to have the strongest impact, followed by protein folding. We speculate that ESM-2 embeddings provide deep contextual understanding of amino-acid sequences due to the model’s large scale training. Protein folding, on the other hand, requires the model to internalize the physico-chemical principles governing that process, which evidently aids in modeling also the protein-ligand interactions. Finally, the scaling part of our attention mechanism is also helpful.

**Table 2: Ablation study results illustrating the effects of different components on performance in the Posebusters benchmark.**
| Ablation | % RMSD $<$ $2$$\mathrm{\SIUnitSymbolAngstrom}$ (↑) | Median RMSD ($\mathrm{\SIUnitSymbolAngstrom}$) (↓) |
| --- | --- | --- |
| RapidDock | 52.1 | 1.90 |
| RapidDock w/o ESM-2 Embeddings | 42.6 | 3.37 |
| RapidDock w/o Folding Pretraining | 46.4 | 2.61 |
| RapidDock w/o Attention Scalers | 48.6 | 2.18 |

## 5 Related Work

Traditionally, blind docking consisted of two stages: pocket finding and molecule pose search. Pocket-finding involves identifying cavities in the protein (Le Guilloux et al., ; Krivák & Hoksza, ; Yan et al., ) while molecule pose search aims to find the optimal molecule conformation in the pocket based on scoring functions
(Koes et al., ; McNutt et al., ).
While still useful, traditional docking approaches are being gradually replaced by
deep learning tools (Corso et al., ; Zhou et al., ; Clyde et al., ).

First successful deep learning models for blind molecular docking applied geometric graph neural networks
(Stärk et al., ; Lu et al., ; Zhang et al., ).
However, their RMSD-minimizing objective has been considered inadequate for modeling the molecule’s position. Instead, DiffDock (Corso et al., ) first formulated molecular docking as generative modeling of the ligand pose distribution, with subsequent approaches adapting that paradigm (Lu et al., ; Nakata et al., ; Corso et al., ).
However, inference involves costly diffusions, making these approaches two-to-three orders of magnitude slower than non-generative methods such as RapidDock.

Recently, a different approach to docking has emerged with RoseTTAFold All-Atom (Krishna et al., ) and AlphaFold-3 (Abramson et al., ).
These tools treat the protein-molecule complex as a single structure to be reconstructed from the protein’s amino acid sequence and the molecule. AlphaFold-3 claims state-of-the-art results in terms of accuracy (76.9% of examples with $\text{RMSD}<2$\mathrm{\SIUnitSymbolAngstrom}$$ on PoseBusters). However, the docking-while-folding framework is extremely computationally expensive. Consequently, AlphaFold-3’s reported inference runtimes are about four orders of magnitude longer than those of RapidDock run on the same hardware. Although incrementally faster docking-while-folding approaches have been published (Qiao et al., ), their speed is still orders of magnitude too slow for allowing large-scale studies envisioned in this paper.

To the best of our knowledge, RapidDock is the first fully transformer-based model for blind molecular docking. Our approach does share, however, several ideas with existing methods. The equivariance of 3D structures representations is a common feature (Stärk et al., ; Lu et al., ; Pei et al., ). Obtaining the molecule position based on predicted distances is done similarly in TANKBind (Lu et al., ) or ArtiDock (Voitsitskyi et al., ). Finally, (Gao et al., ) most recently proposed a symmetry-invariant loss function similar to ours.

## 6 Limitations and Future work

RapidDock achieves excellent accuracy and has a significant speed advantage over other methods. There are several aspects, however, that need further work.

First, we plan to develop a confidence score for
RapidDock’s predictions which is important for decision making.
While this can be done by training another external light model, ideally we would like to extend RapidDock to so it can return its own confidence score.

Second, we plan to fine-tune RapidDock on ligand-protein binding strength prediction, including to identify non-binding ligands. RapidDock’s speed would then allow for efficient identification of all potentially toxic interactions of a drug across the human proteome ($\sim$20k proteins), which is not computationally feasible today and would greatly accelerate the drug discovery process.

Third, we plan to perform more detailed studies of the scaling properties and train larger models on bigger datasets. Among others, we plan to extend pre-training to over 200M protein structures predicted by AlphaFold 2 Varadi et al. (), and use the PLINDER (Durairaj et al., ) protein-ligand datasets for the main training. This will allow us to extend RapidDock to tasks such as protein folding or docking without requiring any input protein structure.

## 7 Conclusions

RapidDock is a truly *fast* and *accurate* molecular docking model that opens new possibilities for in silico drug design. Trained on both protein folding and molecular docking, the model demonstrates a strong understanding of the physicochemical principles behind forming biological structures, achieving accurate results in molecular docking both on holo- and apostructures.

We hope that, in the near future, RapidDock will provide drug designers with a comprehensive view of potential drug interactions across the human proteome. This will unlock important research avenues in biology and machine learning.

Finally, the transformer-based architecture makes the model well-suited for extending pre-training to related biological tasks or fine-tuning on downstream tasks such as predicting toxicity or drug interactions at the cell level, which is one of our current research topics.

#### Acknowledgments

We gratefully acknowledge Adam Dancewicz and Piotr Surma for making this research possible.
We would also like to thank Tadeusz Satława and Karol Musieliński for their insightful comments during the preparation of the manuscript.

## Appendix A Appendix

### A.1 Preprocessing details

##### Matching ligand’s 3D structure with CCD reference

The PDB files contain 3D structures of proteins and molecules. However, only the 3D position and atomic number of a molecule atoms are stored explicitly, while the bonds are not directly available.
We found that methods of inferring molecule graph from a PDB file available in open sources packages lead to incorrect graphs and SMILES (Weininger, ). The correct SMILES is available as CCD code which is also stored in PDB files. A match of the SMILES with the ligand 3D structure from PDB files is performed using Maximum Common Substructure matching (rdFMCS.FindMCS from RDKit package (Bento et al., )).

### A.2 Root Mean Square Deviation (RMSD) Metric

The Root Mean Square Deviation (RMSD) is a standard metric for measuring the average distance between the corresponding atoms in two sets of atomic coordinates. For heavy atoms of the ligand, the RMSD is calculated as:

$$ $\text{RMSD}=\sqrt{\frac{1}{N}\sum_{i=1}^{N}(\mathbf{x}_{i}-\hat{\mathbf{x}}_{i})^{2}},$ (2) $$

where $N$ is the number of heavy atoms, $\hat{\mathbf{x}}_{i}$ is the position vector of the $i$-th heavy atom in the predicted ligand, and $\mathbf{x}_{i}$ is the position vector of the corresponding atom in the ground-truth ligand.

To account for permutation symmetries in the ligand, we use the symmetry-corrected RMSD as described in previous studies (Corso et al., ).

### A.3 Model comparison details

To obtain reliably comparable inference times for the models in [Table 1](#S4.T1), we ran DiffDock-L and NeuralPLexer models on our hardware.
For this, we used inference scripts available at, respectively, [https://github.com/gcorso/DiffDock](https://github.com/gcorso/DiffDock) and [https://github.com/zrqiao/NeuralPLexer/tree/main](https://github.com/zrqiao/NeuralPLexer/tree/main).
The reported metrics come from these evaluation runs. For DiffDock-L, we report metrics on the generated pose with the highest confidence (top-1).

In case of AlphaFold-3, we used the average runtimes reported in Abramson et al. () for the sequence length that was closest to the average sequence length in Posebusters. Because those runtimes were obtained on sixteen A100 GPUs, for a fair comparison, we multiplied them by 16.

### A.4 Training details

The optimal model checkpoints were selected based on the percentage of examples with RMSD $<$ 2Å on the DockGen validation set. The details of the training parameters are shown in [Table 3](#A1.T3).

**Table 3: Training Parameters for Pretraining and Fine-Tuning**
| Parameter | Folding Pretraining | Docking Training |
| --- | --- | --- |
| Dataset Size | $440k$ | $30k$ |
| Masking Factor | $97\%$ | - |
| Task | Amino acid distance prediction | Molecular docking |
| Epochs | $100$ | $200$ |
| Batch Size | $8$ | $4$ |
| Gradient Accumulation Steps | $8$ | $16$ |
| Optimizer | AdamW | AdamW |
| Learning Rate | $1\times 10^{-4}$ | $1\times 10^{-4}$ |
| Max Sequence Length | $1024$ | $2760$ |
| Weight Decay | $0.01$ | $0.01$ |
| Precision | bf$16$ | bf$16$ |
| Training Duration | $48$ hours | $16$ hours |

### A.5 Molecule pose reconstruction details

For solving the optimization problem described in [Section 2.4](#S2.SS4), we use the implementation of the L-BFGS algorithm from PyTorch (torch.optim.LBFGS), so the optimization is also run on the GPU.
We used the following parameters in that function: tolerance_grad $=1$e$-3$, tolerance_change $=$ $1$e$-3$, lr $=1$, max_iter $=100$, max_eval $=500$, history_size $=5$, line_search_fn = strong_wolfe. In particular, using the tolerances of $1$e$-3$ resulted in the same RMSD results on the test sets as with smaller tolerances, but shorter runtimes, on average about two times longer than the transformer forward pass. The optimizer was always able to converge, typically within ten to twenty iterations. We note that although the optimization objective is not smooth, this did not seem to affect the optimizer as we obtained essentially the same results with L1 loss and smooth L1 loss (nn.L1Loss and nn.SmoothL1Loss in PyTorch).

### A.6 Examples of correctly generated poses

While AlphaFold 3 is currently on average more accurate than RapidDock, below we present visualizations of selected examples where RapidDock correctly finds the molecule’s pose while AlphaFold fails to do so.

[Uncaptioned image]: /html/2411.00004/assets/x2.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: * Is there a better way to determine which atoms are rigid and which are flexible? For example, AutoDock Vina determines if bonds are rotatable using a simple chemical definition, which dictates which atoms are flexible and which are not. Just searching through a lot of generated conformations seems like it might miss bonds that only rotate when exposed to external charges.
* Is this model potentially applicable to *target fishing*? Target fishing is the process of taking an existing drug compound and evaluating it against a large number of potential proteins to see if it can target them. This can be applied for drug repurposing, which is the use of currently approved drugs against new indications based on previously unknown binding against a new protein. This is potentially a strong application of the proposed method, but I do not see it explicitly mentioned in the paper anywhere.
soundness: 3
strengths: * Clear presentation of methods and results.
* Novel application of transformer architecture to docking, which results in much faster inference.
* Reasonable design choices in model. These include the addition of features for ligand atom charges and the use of a pre-trained protein language model. The scaling of the attention vector based on distance also seems well-motivated.
* Strong results on two distinct datasets, achieving a better success rate than two competitive deep learning methods at a fraction of the cost. While RapidDock has significantly worse accuracy on Posebusters compared to AlphaFold 3, I do not think this is a negative, because as the authors note AlphaFold 3’s speed is not suitable for large-scale docking. Additionally, my understanding is that AlphaFold-3 performs energy minimization as a post-processing step, which would give it an unfair advantage compared to RapidDock.
* Ablations of various model components help show the benefits of each design choice.
summary: The authors tackle the problem of proteome-scale docking, the goal of which is predicting the binding pose of a ligand against many thousands of proteins. To do this, they develop an equivariant Transformer model (RapidDock) for dramatically accelerating docking compared to previous diffusion or GNN-based approaches. The model takes various features from the protein and ligand as input, and outputs a prediction of the binding pose of the ligand. The results show that RapidDock achieves 100x faster runtimes than three competing deep learning methods, while retaining equivalent accuracy (except to AlphaFold 3, which is much slower but has much better accuracy).
weaknesses: * Motivation behind the problem setting is unclear. The authors address proteome-scale docking because any protein in the human proteome could be a potential *off-target* (a term of art that should probably be included in the paper) of a drug. Thus, docking against all proteins and then predicting affinity in a downstream task would detect these potential off-targets before they are discovered in later preclinical or clinical testing. However, I am not convinced that docking to each protein in the proteome is necessary to detect off-target effects. Pharmacologists often screen a drug against a limited number of safety-relevant proteins (in the hundreds), such as G-protein coupled receptors or ion channels, that are frequent off-targets [1, 2]. This is usually sufficient to detect many clinical issues beforehand, and it is not obvious that just considering more potential off-targets would further reduce the rate of adverse effects occurring (which for example may be due to more complex issues, such as toxicity of metabolic products or on-target unwanted effects). Could the authors provide a better argument for why it is important to screen a drug against all potential protein targets?
[1] Bendels et al. “Safety screening in early drug discovery: An optimized assay panel.“ J Pharmacol Toxicol Methods 2019.
[2] Peters et al. “Can we discover pharmacological promiscuity early in the drug discovery process?” Drug Discovery Today 2012.

* Limited baseline comparisons. I think the most important baseline to include would be AutoDock Vina (or another similar docking program). Despite not being deep learning-based, Vina is relatively fast and the current state-of-the-art in applied fields. Practitioners conducting large-scale docking would likely use Vina, so including results on this baseline is important. An additional deep-learning based docking model, such as TANKBind, would also improve the strength of the results, but it is not as critical.

* Some of the design choices are not well-explained. For example, it is not clear to me why discretized charge embeddings are used for the ligand atoms instead of simply providing the charge scalar as an input. It is also not clear what role the distance bias matrices play.

---

confidence: 3
contribution: 3
initial_rating: 3
presentation: 1
questions: 1. DiffDock-L (also AlphaFold3 and NeuralPLexer) is a generative approach, so it gives multiple binding poses according to their probabilities. Therefore, its prediction accuracy can be improved by including multiple different poses (Top-n poses) in the evaluation. However, RapidDock is deterministic, so it gives only a single output. It is understood that the authors want to emphasize computational speeds, but it seems they also need to discuss the accuracy aspects for a fair comparison.
2. Since RapidDock requires conformation searches for each molecule before docking, the authors need to clarify whether the reported computational times include the time for conformer searches or not. 
3. The proposed method has been compared with AlphaFold3 and NeuralPLexer. However, the former performs a rigid docking, whereas the latter predicts binding poses only from protein sequences and molecular graphs. Therefore, the prediction complexity of RapidDock is much lower than that of the baseline models, so the direct comparison between them is less meaningful.  The authors need to clarify this fact in the introduction or result sections. The current form may cause undesirable confusion to potential readers.
4. In the reconstruction of ligand location, predicted distances between ligands and also between ligand-protein may not precisely match with the coordinates of a single pose. If this is the case, the authors should elaborate more details about this process. Does it require a kind of post-processing?
5. The authors greatly emphasize the importance of proteome-wide docking, but this work does not provide any meaningful analysis except the computational time, which can be readily estimated without performing the actual calculations. The authors may add an additional study verifying that the proposed method can indeed provide meaningful results from the proteome-wide docking. Otherwise, they need to tone down their argument from the title and introduction parts. 
6. Appendix A.6 shows the examples of 3D structures predicted by RapidDock and AlphaFold3. The authors deliberately selected specific examples where RapidDock outperformed AlphaFold3, while they admit that the latter is far better than the former on average. AlphaFold3 even predicted those structures from the sequence-level information. These examples may lead to misunderstanding that RapidDock works better than AlphaFold3. The authors need to provide examples where RapidDock fails while AlphaFold3 succeeds.
soundness: 1
strengths: 1. This work shows the possibility of transformer-based approaches for binding structure predictions, while most previous works are based on diffusion methods.
2. The proposed method outperformed the popular base-line, DiffDock-L, in the two benchmark studies, while its computational time for predictions is much faster than those of all base-line models.
summary: This work proposed a fast and reliable prediction method, RapidDock, for protein-ligand binding poses. Most previous methods employed in this work as base-lines used diffusion models, leading to large computational costs. However, RapidDock is based on a transformer and so much faster than the base-line models. The benchmark studies on PoseBusters and DockGen show that RapidDock is not only fast but also far more accurate than others except AlphaFold3. While the performance of the proposed method seems competitive, there are several issues that need to be addressed.
weaknesses: 1. The proposed method needs to generate 96 molecular conformations for each molecule and analyze the conformers to obtain its distance matrix.
2. The experiment in Section 4.2 seems meaningless because it uses holostructures when it predicts binding poses.
3. The title and introduction parts emphasize the importance of proteome-wide docking, but this work does not provide any meaning results regarding that. 
4. Technical details of the proposed method are insufficient.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 2
questions: - When comparing runtime in inference mode, did you preprocess the protein, or was the comparison done solely for molecule conformation?
- Did you use the same time-splitting approach as previous methods, such as DiffDock and NeuralPlexer, for the dataset?
- In the statement, 
> "Only the fixed distances across the molecule’s possible conformations are recorded, and others are denoted by a special value of −1," 

  Could you clarify why you selected -1 as the special value?
- In the RapidDock attention section, you state, 
> "First, we multiply the attention scores corresponding to input pairs with known distances (i.e., ligand-ligand within a rigid part and protein-protein) by a learnable scalar $s_m$, one for each layer $m$." 
  
  Did you also apply this to protein-ligand pairs? If not, could you explain why?
- For inference, do you generate one conformation per runtime for each protein-ligand pair?
soundness: 3
strengths: - The authors demonstrate the need for a faster model by outlining the scalability limitations of previous deep learning models.
- Aligned with their motivation, RapidDock shows the ability to perform conformation sampling for molecular docking in GPU inference runtime in approximately one-hundredth of a second per protein-ligand pair.
- In benchmarking with PoseBuster, RapidDock achieves the best performance among open-source codes (noting that AlphaFold 3 is not open source), particularly in the percentage of ligands achieving RMSD < 2 Å.
- The paper provides a clear explanation of how ligand and protein modalities are utilized and fed into the Transformer model.
- For constructing the ligand distance matrix, the authors use a hybrid approach, incorporating physics-based methods from RDKit, such as MMFF and EDKG.
- Additionally, the inclusion of ligand charge embeddings represents another hybrid approach in the model.
- For protein embeddings, the authors showcase the effectiveness of using not only pre-trained models but also their custom-trained ESM-2 models.
- The Transformer architecture employs a non-autoregressive approach with a full attention mask, introducing a new method for molecular docking by incorporating an attention scaler within the attention mechanism.
- The training hyperparameters are shared in detail through comprehensive tables.

**Originality:** RapidDock introduces a unique approach to molecular docking, particularly in terms of preprocessing compared to other DL-based methods. The detailed steps in the ligand and protein embedding process highlight its originality, which the authors further validate through an ablation study.

**Significance:** From a large-scale proteomic perspective, RapidDock is highly scalable and significantly faster in runtime compared to other DL-based and search-based molecular docking methods.
summary: In this paper, the authors introduce RapidDock, a new approach to molecular docking that leverages a Transformer model. The method includes both ligand atom embeddings and ligand charge embeddings. Protein representations are generated using embeddings from protein amino acids, the ESM-2 PLM, and calculated distance metrics. Instead of deep learning, RDKit-based methods such as MMFF and EDKG, construct the rigid distance matrix for molecules. Trained on the PDBBind and BindingMOAD datasets, RapidDock outperforms DiffDock and other open-source methods on the PoseBuster benchmark in % of ligands in RMSD < 2 Å metric, delivering at least a 100x increase in inference speed.
weaknesses: - The code is not shared.
- Although the method section claims equivariance, it lacks sufficient explanation on this aspect.
- The rationale for using ligand atom charges is not adequately clarified.
- It is unclear why non-fixed distances in the molecule's rigid distance matrix are assigned a value of -1.
- The annotations for distance bias matrices are insufficiently explained; the annotations appear to be included simply because they work, without detailing why they are effective.
- Similarly, the rationale behind RapidDock’s use of attention and charge embeddings, along with their annotations, is not fully addressed.
- The splitting strategy for the training, validation, and test sets is not sufficiently described.
- The parameter comparison in benchmarking does not seem fair; DiffDock results are compared with 30 million parameters, while RapidDock has 60 million parameters.
- The RMSD metric, commonly used in molecular docking and Structure-Based Drug Design (SBDD), does not always yield bioactively, physically, or chemically plausible structures, as shown by Posebuster[1], PoseCheck[2], PoseBech[3], and CompassDock[4]. Including these metrics in the benchmark would strengthen the study.
- Appendix A.6 lacks a comparison between DiffDock and NeuralPLexer examples.
- The extent of ligand filtering during ligand preparation is not sufficiently discussed.

**Quality:** Although the authors claim this is the first use of a Transformer-based model in blind docking, ETDock[5] and FeatureDock[6] have previously used Transformers for molecular docking; however, these methods are not mentioned in the paper. Additionally, the benchmarking is limited to comparisons with only a few popular methods. In the conclusion, the authors state that: 

>"... the model demonstrates a strong understanding of the physicochemical principles behind forming biological structures," 

yet no bioactivity or physicochemical analyses, as discussed earlier, have been conducted to support this claim.

**Clarity:** The paper contains numerous grammatical errors that detract from readability and should be carefully revised. Placing the Related Work section after the Experiments section disrupts the flow, and the method annotations are not clearly explained.

**Reproducibility:** As the code is not shared, it is currently impossible to test whether it performs as reported. If the code were provided, I would be able to review, test, and reassess my evaluation (including scoring and comments) accordingly.

### **References**
[1] Martin Buttenschoen, Garrett M Morris, and Charlotte M Deane.  Posebusters:  Ai-based docking methods fail to generate physically valid poses or generalise to novel sequences. Chemical Science, 15(9):3130–3139, 2024. 

[2] Charles Harris, Kieran Didi, Arian R Jamasb, Chaitanya K Joshi, Simon V Mathis, Pietro Lio, and Tom Blundell. Benchmarking generated poses: How rational is structure-based drug design with generative models? arXiv preprint arXiv:2308.07413, 2023.

[3] Alex Morehead, Nabin Giri, Jian Liu, Jianlin Cheng. Deep Learning for Protein-Ligand Docking: Are We There Yet? arXiv preprint arXiv:2405.14108, 2024.

[4] Ahmet Sarigun, Vedran Franke, Bora Uyar, Altuna Akalin. CompassDock: Comprehensive Accurate Assessment Approach for Deep Learning-Based Molecular Docking in Inference and Fine-Tuning. arXiv:2406.06841, 2024.

[5] Yiqiang Yi, Xu Wan, Yatao Bian, Le Ou-Yang, Peilin Zhao. ETDock: A Novel Equivariant Transformer for Protein-Ligand Docking. arXiv:2310.08061, 2023.

[6] Mingyi Xue, Bojun Liu, Siqin Cao, Xuhui Huang. FeatureDock: Protein-Ligand Docking Guided by Physicochemical Feature-Based Local Environment Learning using Transformer. ChemRxiv

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: Wording:
L. 30 “ … will revolutionize medicine” is an unsupported statement, please reformulate. It would be better to focus on the obstacles that need to be overcome and how you tackle those (like in
your paragraph 2, L. 34 X.).
In l.113f. the authors claim equivariance, however they work with interatomic distances only, making your model invariant.
Permutation loss citation is wrong (e.g. l. 228), cite original paper (Zhu et al. 2022 Direct molecular conformation generation) 
In my opinion, l.419f. “the model demonstrates a strong understanding of the physicochemical principles” is a big stretch - how is this supported in your work? Accurate and fast prediction of ligand poses in proteins doesn’t mean that the model understands the physicochemical principles that lead to those poses, and it is debatable whether this is even needed.
Open questions:
Since the prediction time is critical to the paper's claims, I would like to see how much time you spend at inference on average in a) pre-processing, b) prediction, c) reconstruction and d) any form of post-processing.

In l.162f., what is the reasoning for choosing 257 buckets? Are there ablations on higher/lower resolution? How does this affect the performance-inference time tradeoffs? Same questions for the charge embeddings.
soundness: 3
strengths: RapidDock is fast and accurate, allowing it to tackle human proteome-scale docking studies. The reported speed could enable its use as an oracle function for other DD-related tasks in future work. Ablations on PoseBusters and Dockgen show promising results. The
paper is well written and gives a lot of insights into the modelling and training
summary: The paper introduces RapidDock, a fast and accurate transformer-based model for the
blind-docking task. The model predicts interatomic distances. Afterwards, the docked
pose is reconstructed with the L-BFGS algorithm. The authors report a 100x speed-up while
simultaneously improving over commonly used deep learning-based docking models in the
PoseBusters and DockGen datasets. The authors make first experiments at human
proteome-scale docking.
weaknesses: While the method shows good results on a task relevant to computational biology, there is not enough novelty on the ML side that justifies acceptance as a main track contribution. The authors use a standard transformer embedding the ligand and proteins with ESM2 and
employ cross-attention to the distance matrix. I encourage submission to a domain-specific journal.

Meta-review: The paper proposes a rapid docking method based on a Transformer. The model takes as input a ligand and a protein, both of which are tokenized, and their 3D structures are passed through all self-attention layers. The protein is embedded using ESM-2.

The method achieves strong empirical results, and the reviewers appreciated this aspect. Overall, the work is very well executed. The training choices, such as the use of AlphaFold2-generated structures or biasing self-attention with distance, are generally intuitive. (Additional ablations would be very useful.)

However, three out of four reviewers voted for rejection.

A significant drawback of the manuscript is its limited novelty. The idea of using deep learning models to directly predict poses and/or docking scores has been explored before, as noted by one of the reviewers. UniMol2 is one such comparison point.

Another major critique focused on the limited evaluation. More detailed analysis of poses or comparisons to more docking baselines were suggested. One reviewer specifically proposed using a time-based split. It would also be useful to consider other routes to accelerating existing docking methods (e.g., distilling DiffDock).

These lower-level comments on the evaluation scheme reflect a broader issue: the paper does not clearly demonstrate its intended use. Two reviewers expressed doubts about the possibility of proteome-wide docking. For example, it is not fully clear how broadly the method generalizes to novel ligand structures or different protein families, which might be crucial depending on the use case. The paper would benefit from a clearer explanation of the model’s practical applications, ideally supported by an actual demonstration.

All in all, the paper fell slightly short of the acceptance threshold. While its strong empirical results were acknowledged, it did not fully clarify how the model would be used in practice. At this stage, I am recommending rejection. I hope these comments will be helpful in improving the paper.

Decision: Reject