High performance binding affinity prediction with a Transformer-based surrogate model

In the current paradigm of drug discovery pipelines, identification of compounds that bind to a target with high affinity constitutes the first step. This is typically performed using resource-intensive experimental methods to screen vast chemical search spaces - a key bottleneck in the drug-discovery pipeline. To streamline this process, highly-scalable computational screening methods with acceptable fidelity are needed to screen larger portions of the chemical search space and identify promising candidates to he validated using experiments. Machine learning methods, namely, surrogate models have recently evolved into favorable alternatives to perform this computational screening. In this work, we present Simple SMILES Transformer (SST), an accurate and highly-scalable binding affinity prediction method that approximates the computationally-intensive molecular docking process using an encoder-only Transformer architecture. We benchmark our model against two baselines that feature fundamentally different approaches to docking surrogates: RegGO, a MORDRED fingerprint based multi-layer perceptron model, and Chemprop, a directed message-passing graph neural network. Unlike Chemprop and RegGO, our method operates solely on the SMILES representation of molecules without needing additional featurization, which leads to reduced preprocessing overhead, higher inference throughput and thus better scalability. We train SST in a distributed fashion on the Polaris supercomputer at the Argonne Leadership Computing Facility (ALCF). We then deploy it at an unprecedented scale for inference across 256 compute nodes of ALCF's Aurora supercomputer to screen 22 billion compounds in 40 minutes in search of hits with high binding affinity to oncoprotein RtcB ligase. SST predictions emphasize several molecular motifs that have previously been confirmed to interact with residues in their target binding pockets.