Learning Compressed Shape-Aware Molecular Representations for Virtual Screening

Virtual screening of billion-scale molecular libraries based on 3D shape similarity remains computationally prohibitive, requiring expensive conformational sampling and alignment, as done by established tools like *ROCS*. Here, we introduce *SAND* (**S**hape-**A**ware **N**eural **D**escriptor), a method that can retrieve shape similar molecules from their 2D graph alone. Our approach makes two key contributions: (1) a rank-preserving contrastive learning framework using differentiable Spearman correlation that results into representations where similarity strongly correlates with 3D molecular shape overlap (R=0.86), and (2) an end-to-end learned quantization-aware training scheme that jointly optimizes the encoder with a two-level IVF-PQ discretization step, achieving approximately $4\times$ better compression than post-hoc quantization at equivalent retrieval quality. We demonstrate that *SAND* enables searching over 10 billion molecules in less than a second on a single GPU node - a speedup of $>10^{8} \times$ compared to traditional methods. We release open-source code and trained weights to facilitate adoption.

Finding new medicines often starts with searching huge digital libraries of molecules — sometimes billions of them — to find ones that have a similar 3D shape to a known drug. Shape matters because molecules that "fit" the same biological target tend to look similar in three dimensions. The problem is that comparing 3D shapes is very slow: every molecule can fold into many possible 3D arrangements, and each one must be rotated and matched against the query. With today's libraries growing past ten billion compounds, this kind of search has become practically impossible. We trained a neural network called SAND that learns to predict 3D shape similarity directly from a molecule's flat 2D chemical structure, skipping the expensive 3D calculations entirely. We also taught it to compress each molecule into a small digital fingerprint that can be searched almost instantly. The result: searching ten billion molecules now takes less than a second on a single computer — over a hundred million times faster than traditional tools — making large-scale shape-based drug discovery more accessible.