Learning Compressed Shape-Aware Molecular Representations for Virtual Screening

Robin Winter, Julian Cremer, Djork-Arné Clevert
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:135187-135199, 2026.

Abstract

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 (Shape-Aware Neural Descriptor), 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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-winter26a, title = {Learning Compressed Shape-Aware Molecular Representations for Virtual Screening}, author = {Winter, Robin and Cremer, Julian and Clevert, Djork-Arn\'{e}}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {135187--135199}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/winter26a/winter26a.pdf}, url = {https://proceedings.mlr.press/v306/winter26a.html}, abstract = {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 (Shape-Aware Neural Descriptor), 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.} }
Endnote
%0 Conference Paper %T Learning Compressed Shape-Aware Molecular Representations for Virtual Screening %A Robin Winter %A Julian Cremer %A Djork-Arné Clevert %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-winter26a %I PMLR %P 135187--135199 %U https://proceedings.mlr.press/v306/winter26a.html %V 306 %X 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 (Shape-Aware Neural Descriptor), 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.
APA
Winter, R., Cremer, J. & Clevert, D.. (2026). Learning Compressed Shape-Aware Molecular Representations for Virtual Screening. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:135187-135199 Available from https://proceedings.mlr.press/v306/winter26a.html.

Related Material