Gpu acceleration enables fast large scale drug molecule searches
SweepANN: A GPU-Accelerated Large-K Approximate Vector Search for Drug Discovery
DatabasesDistributed, Parallel, and Cluster ComputingInformation Retrieval
Summary
Finding new drug candidates requires searching through vast libraries of molecules to see which might fit well into protein pockets. The authors improved this search by using multiple GPUs to quickly scan and find many top matching molecules for each pocket with very high accuracy. Their method, SweepANN, speeds up searches dramatically compared to previous approaches. This lets drug discovery efforts handle hundreds of millions of molecules in minutes rather than hours or days.
What this means in practice
- •For pharmaceutical researchers: Enable large-scale candidate molecule retrieval in drug design platforms, reducing search times from hours to minutes while maintaining near-perfect recall.
- •For machine learning engineers: Implement GPU-accelerated vector search in systems requiring fast, large-scale approximate nearest neighbor queries with very large top-K outputs.
Authors
Zhenyu Zhang, Qi Lin, Nan Zhang, Zhichao Cao
Abstract
DrugCLIP places small molecules and protein pockets in one space, turning genome-wide drug discovery into large-K retrieval by vector proximity. A campaign fixes over 20K pockets in advance for one library of over 500M molecules. Each pocket keeps K = 500K candidates, 0.1% of the collection, versus 100 in ANN benchmarks. Missed hits never get tested, so mean recall must reach 99% relative to exact search. Large answers make IVF probe every list holding a result and graph search track K or more candidates, so costs grow with K. On our CPUs and GPUs, neither index family reaches the target without losing most of its edge over scanning. SweepANN partitions compressed vectors across GPUs and scans each resident partition. One matrix multiplication scores each block for a whole query batch, reading it once. Several plans perform the K-dependent selection and merging of each query's K answers, each returning an exact top-K under compressed scoring. The window plan retains each query's best K from sampled blocks as its seed, and only rows above each seed's minimum score, shared once across GPUs, enter the buffers. Four Blackwell GPUs (RTX PRO 6000) answer all 26,562 pockets with a resident time of about 61 seconds, excluding library loading, and 99.6% mean recall relative to exact FP32 search. DrugCLIP reports roughly 24 hours with eight A100 GPUs in its own campaign. On 500M rows, the window plan runs 1.41x faster than SweepANN's grouped plan in the same memory. Across 25M-500M rows and K of 10K-500K, all GPU IVF configurations meeting the recall target take 8.2-24x SweepANN's resident time. On DEEP1B's 1B vectors across eight GPUs, against SweepANN's exact full-score scan of uncompressed FP32 rows, the window plan cuts resident time by 5.4x and 3.2x at K = 100K and 500K. It also cuts per-GPU peak memory by 6.0x and 5.3x, losing 1.0 and 0.7 mean-recall points.