Drug Discovery under Covariate Shift with Domain-Informed Prior Distributions over Functions
Leo Klarner, Tim G. J. Rudner, Michael Reutlinger, Torsten Schindler, Garrett M. Morris, Charlotte M. Deane, Yee Whye Teh
Abstract
Accelerating the discovery of novel and more effective therapeutics is an important pharmaceutical problem in which deep learning is playing an increasingly significant role. However, real-world drug discovery tasks are often characterized by a scarcity of labeled data and significant covariate shifta setting that poses a challenge to standard deep learning methods. In this paper, we present Q-SAVI, a probabilistic model able to address these challenges by encoding explicit prior knowledge of the data-generating process into a prior distribution over functions, presenting researchers with a transparent and probabilistically principled way to encode data-driven modeling preferences. Building on a novel, gold-standard bioactivity dataset that facilitates a meaningful comparison of models in an extrapolative regime, we explore different approaches to induce data shift and construct a challenging evaluation setup. We then demonstrate that using Q-SAVI to integrate contextualized prior knowledge of drug-like chemical space into the modeling process affords substantial gains in predictive accuracy and calibration, outperforming a broad range of state-of-the-art self-supervised pre-training and domain adaptation techniques.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers4
- Context-Guided Diffusion for Out-of-Distribution Molecular and Protein DesignLeo Klarner, Tim G. J. Rudner, Garrett M. Morris, Charlotte M. Deane et al.ICML 2024 · 18 citations
- Should We Learn Most Likely Functions or Parameters?Shikai Qiu, Tim G. J. Rudner, Sanyam Kapoor, Andrew Gordon WilsonNeurIPS 2023 · 15 citations
- Instructor-inspired Machine Learning for Robust Molecular Property PredictionFang Wu, Shuting Jin, Siyuan Li, Stan Z. LiNeurIPS 2024 · 14 citations
- REACT: Residual-Adaptive Contextual Tuning for Fast Model Adaptation in Threat DetectionJiayun Zhang, Junshen Xu, Bugra Can, Yi FanWWW 2025 · 3 citations
Builds on6
- WILDS: A Benchmark of in-the-Wild Distribution ShiftsPang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie et al.ICML 2021 · 1,773 citations
- Strategies for Pre-training Graph Neural NetworksWeihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik et al.ICLR 2020 · 1,744 citations
- E(n) Equivariant Graph Neural NetworksVictor Garcia Satorras, Emiel Hoogeboom, Max WellingICML 2021 · 1,432 citations
- Self-Supervised Graph Transformer on Large-Scale Molecular DataYu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie et al.NeurIPS 2020 · 1,113 citations
- Tractable Function-Space Variational Inference in Bayesian Neural NetworksTim G. J. Rudner, Zonghao Chen, Yee Whye Teh, Yarin GalNeurIPS 2022 · 70 citations
Related papers
- Test-Time Adaptation without Source Data for Out-of-Domain Bioactivity PredictionYiming Yang, Zhiyuan Zhou, Yueming Yin, Hoi-Yeung Li et al.ICLR 2026
- 3D Infomax improves GNNs for Molecular Property PredictionHannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou et al.ICML 2022 · 269 citations
- 3DLinker: An E(3) Equivariant Variational Autoencoder for Molecular Linker DesignYinan Huang, Xingang Peng, Jianzhu Ma, Muhan ZhangICML 2022 · 65 citations
- CoDrug: Conformal Drug Property Prediction with Density Estimation under Covariate ShiftSiddhartha Laghuvarapu, Zhen Lin, Jimeng SunNeurIPS 2023 · 14 citations
- Learning Invariant Molecular Representation in Latent Discrete SpaceXiang Zhuang, Qiang Zhang, Keyan Ding, Yatao Bian et al.NeurIPS 2023 · 41 citations
