APEX: Approximate-but-exhaustive search for ultra-large combinatorial synthesis libraries

AI summary

Overview

Research area: Machine learning for drug discovery, specifically virtual screening over ultra-large combinatorial synthesis libraries (CSLs), with a focus on surrogate modeling, factorization-based retrieval, and GPU-accelerated top-k search.

Technical level: Advanced. The paper assumes familiarity with molecular representations (SMILES, 2D graphs), neural network surrogates, docking scoring functions, Bayesian optimization, and GPU kernel design. The core mathematical development (surrogate parameterization, factorized embeddings, associative contributions) is dense.

Scope: The paper presents APEX, a search protocol that combines a neural network surrogate with a learned reaction factorizer so that the entire enumerated chemical space of a CSL (tens of billions of compounds) can be evaluated exhaustively on a single GPU in seconds, and benchmarks it against ground-truth docking scores on five targets plus a comparison to Thompson sampling.

What This Paper Is About

Make-on-demand CSLs such as Enamine REAL contain tens of billions of compounds, but scoring functions like docking are too expensive to evaluate across the whole library, so practical screening campaigns typically score fewer than 0.1% of available compounds and can miss high-scoring molecules. APEX attacks this by learning to decompose a surrogate model's prediction for any library compound into a sum of pre-cached, per-synthon scalar contributions, so that every compound in the library can be scored with a few floating-point operations and the exact top-k set under the approximate model can be retrieved in one exhaustive GPU pass. The goal is to make declarative top-k queries—maximize an objective subject to physicochemical constraints—fast and amortized, without retraining when the objective or constraints change.

Key Contributions

1. The APEX factorization. A decomposition in which a molecule's surrogate prediction is written as the sum of synthon associative contributions plus a bias, f̂̂ᵢ(x) = Σ₍ᵣ,ₛ₎∈χ vᵢ,ᵣ,ₛ + bᵢ, reducing exhaustive evaluation over the CSL to a cached lookup over synthon vectors. A noise-injection term during surrogate training (isotropic normal noise on embeddings) statistically regularizes the surrogate against the errors-in-variables problem introduced by this later factorization.

2. A new benchmark CSL. An open, reproducible CSL built by sampling 1 million "lead-like" compounds from ZINC22 and fragmenting them with BRICS, then enumerating two- and three-component reactions. The final library contains over 10B molecules, with smaller downsampled versions of over 12M and 1M products that are fully enumerated so ground-truth top-k sets can be computed. Synthetic feasibility was not considered or ensured.

3. An annotated benchmark dataset. All molecules in the smaller enumerated library were embedded with RDKit and docked against five targets (PARP1, MET, DRD2, F10, ESR1) using a GPU implementation of the AutoDock Vina scoring function, with receptor structures and binding sites from DOCKSTRING, plus RDKit physicochemical properties.

4. Custom CPU and GPU top-k implementations. A PyTorch CUDA C++ extension module using a chain-of-batches strategy and the GPU-compatible AIR top-k algorithm, with a priority queue of size k ordered lexicographically on constraint violation and objective value. Code is released at https://github.com/NumerionLabs/apex.

Main Findings

• High recall at modest budget: With k = 100,000 retrievals (0.803% of compounds in the 12M CSL), APEX recovered the ground-truth top-j compounds at rates far exceeding a random baseline across all five targets and all constraint sets; a random baseline achieves recall below 0.01.

• Constraint handling works: Constraint satisfaction rates for APEX retrievals exceeded the base fraction of satisfying compounds in the library for every constraint set tested. Recall was highest for the Veber constraints (least stringent, satisfied by most compounds); the Astex Rule-of-3 set was the most stringent and had the lowest satisfaction rate, but still above the library baseline.

• Outperforms or matches Thompson sampling: Against the Thompson sampling implementation of Klarich et al. (2024), evaluated on the top five largest reactions in the 12M CSL (over 4 million products in total, run for 100, 1,000, or 10,000 iterations with 3 warmup steps for two-component and 10 for three-component reactions), APEX consistently outperformed or matched TS at recalling top-j compounds, with particular strength at lower evaluation budgets. The comparison was unconstrained only, because that TS implementation does not directly support constraints on molecular properties.

• Docking score enrichment: The empirical CDF of docking scores for APEX top-k sets shows clear enrichment over the background distribution from the 12M library, in both the 10B and 12M CSLs.

• Zero-shot transfer to Enamine REAL: Applied without retraining to the Enamine REAL library (9-2024 version, more than 70B compounds), APEX enriched docking scores relative to a background of 100,000 random compounds and, in most cases, relative to the top-k of the 10B BRICS library. The lowest enrichment was on MET kinase, which also corresponded to the largest drop-off in R-squared for the factorizer.

• Runtime: With GPU top-k, unconstrained retrieval of k = 1,000,000 from the 10B BRICS library took 10.9 seconds averaged over five runs, and with Lipinski Rule of 5 constraints 10.8 seconds; the paper describes retrieving k = 1 million from a 10 billion compound CSL in approximately 30 seconds on a single T4 GPU. For the 70B Enamine REAL library, unconstrained k = 1,000,000 took 184.0 seconds on GPU versus 858.9 seconds on CPU; with Lipinski Rule of 5 it took 202.4 seconds on GPU versus 3,142.5 seconds on CPU. CPU runtime increased approximately linearly with the number of constraints added, while GPU runtime did not.

• Theoretical scaling example: For a CSL with one three-component reaction and 10,000 synthons per R-group (|S| = 30,000, one trillion products), APEX requires 30,000 neural network evaluations instead of one trillion; with d = 1024 the cached associative embedding matrix is about 120 MB versus about 4 PB for latent representations of all products. The dot products require 2|S|×d − |S| = 61.41 million floating-point operations, and evaluating all one trillion compounds requires roughly 3 trillion floating-point operations, versus the 8.1 TFLOPS capability of an NVIDIA Tesla T4.

• Model configuration: The surrogate is trained only on the 1M compound CSL; the factorizer is trained on either the 12M or 10B CSL without labels. Embedding dimension d = 64, and no extensive hyperparameter tuning was performed.

Methodology in Plain English

The approach has three stages. First, train a surrogate. A neural network encoder maps a molecule to a d-dimensional embedding, and each prediction task (the docking objective or a constraint) is modeled as a linear function of that embedding, wᵢᵀgθ(x) + bᵢ. The surrogate is trained on the 1M labeled CSL, with noise added to embeddings during training so the later factorization stays robust.

Second, train a factorizer. Because CSLs are built hierarchically—reactions contain R-groups, R-groups contain synthons—the factorizer learns an encoder at each level (SynthonEncoder, RgroupEncoder, ReactionEncoder) and then two additional encoders (SynthonValueEncoder, RgroupKeyEncoder) that produce a value vector per synthon and a key matrix per R-group. Multiplying the key by the value yields a per-synthon contribution, and the molecule's embedding is reconstructed as the sum of these contributions over its (R-group, synthon) assignments. The factorizer is trained to minimize the reconstruction error of the frozen surrogate's embeddings, using no labels.

Third, run the search. Once trained, the per-synthon contributions for a given task can be precomputed and cached. Any query—objective plus constraint bounds—reduces to summing a handful of cached scalars per compound and adding a bias. All compounds are streamed through, each entered into a priority queue of size k ordered lexicographically on a combined constraint-violation and objective vector, and compounds with any constraint violation are removed at the end. Because the surrogate is approximate, the retrieved set is exact for the approximate model but only approximate with respect to the true objective.

Why This Matters

Impact on research: APEX removes a structural bottleneck in virtual screening. Instead of treating the CSL as a haystack to be sampled, it makes the whole haystack evaluable, and it separates one-time training from query execution, so changing an objective or constraint requires no retraining or policy re-optimization. The paper also contributes an open benchmark CSL with ground-truth docking scores, which allows any virtual screening algorithm to be compared against known top-k sets.

Real-world applications:

• Hit identification in drug discovery projects, where teams need to rank candidate compounds against a target under drug-likeness filters.
• Rapid re-querying during a campaign as project constraints tighten or change (for example, swapping a property cutoff), without new oracle calls.
• Screening of commercial vendor libraries such as Enamine REAL in a zero-shot manner before any project-specific data is collected.
• Generating focused, diverse shortlists of compounds for procurement and experimental testing at scales (hundreds of thousands to a million compounds) that iterative methods reach only slowly.

Industry relevance: The paper is a collaboration in which four authors were employed by Numerion Labs, which develops and uses APEX in its virtual screening platform, and four are employed by NVIDIA, which develops CUDA (in which the factorized top-k is implemented) and the GPUs used in experiments. The runtime profile—retrieving one million compounds from a 70B vendor library in roughly three minutes on a single consumer/server-class T4 GPU—is directly relevant to industrial screening throughput and cloud cost.

Future Directions

• Target-specific fine-tuning. The authors state that greater docking score enrichment than the zero-shot Enamine REAL results is likely achievable by fine-tuning the surrogate, and subsequently the factorizer, using labeled data from the target CSL. The MET kinase result, with the largest R-squared drop-off, is the clearest candidate for this treatment.

• Combining factorization with protein–ligand co-embedding models. The paper suggests the factorization strategy could in principle be applied on top of methods like ConPLex, DrugCLIP, and LigUnity, which currently require encoding library compounds individually (DrugCLIP's pre-encoding of a 6-billion compound library is cited at on the order of 30 hours per target).

• Broader and harder constraint regimes. The current evaluation covers Veber, Lipinski Rule of 5, and Astex Rule-of-3 constraint sets; the authors note that constraint handling is a key differentiator from iterative methods such as V-SYNTHES and Thompson sampling, which do not natively handle property constraints.

• Synthesizability and library realism. The benchmark CSL was built by BRICS fragmentation and explicitly does not consider or ensure synthetic feasibility, so transferring these results to experimentally realizable libraries remains an open question. The paper does not report a formal limitations section, nor extensive hyperparameter studies for the surrogate or factorizer beyond the stated d = 64 configuration.

Target Audience

Computational chemists and cheminformatics practitioners running virtual screening campaigns at scale; machine learning researchers working on retrieval, surrogate modeling, or factorization methods over structured or compositional data; and drug discovery teams at pharmaceutical and biotech companies who need fast, constraint-aware shortlisting over vendor or in-house libraries. Readers without a background in molecular representation learning, docking, or GPU programming will find the methodology sections demanding, though the problem framing and reported runtimes are accessible.

Authors’ abstract

Make-on-demand combinatorial synthesis libraries (CSLs) like Enamine REAL have significantly enabled drug discovery efforts. However, their large size presents a challenge for virtual screening, where the goal is to identify the top compounds in a library according to a computational objective (e.g., optimizing docking score) subject to computational constraints under a limited computational budget. For current library sizes -- numbering in the tens of billions of compounds -- and scoring functions of interest, a routine virtual screening campaign may be limited to scoring fewer than 0.1% of the available compounds, leaving potentially many high scoring compounds undiscovered. Furthermore, as constraints (and sometimes objectives) change during the course of a virtual screening campaign, existing virtual screening algorithms typically offer little room for amortization. We propose the approximate-but-exhaustive search protocol for CSLs, or APEX. APEX utilizes a neural network surrogate that exploits the structure of CSLs in the prediction of objectives and constraints to make full enumeration on a consumer GPU possible in under a minute, allowing for exact retrieval of approximate top-k sets. To demonstrate APEX's capabilities, we develop a benchmark CSL comprised of more than 10 million compounds, all of which have been annotated with their docking scores on five medically relevant targets along with physicohemical properties measured with RDKit such that, for any objective and set of constraints, the ground truth top-k compounds can be identified and compared against the retrievals from any virtual screening algorithm. We show APEX's consistently strong performance both in retrieval accuracy and runtime compared to alternative methods.