Drug-TTA: Test-Time Adaptation for Drug Virtual Screening via Multi-task Meta-Auxiliary Learning

Ao Shen, Mingzhi Yuan, Yingfan MA, Jie Du, Qiao Huang, Manning Wang

Abstract: Virtual screening is a critical step in drug discovery, aiming at identifying potential drugs that bind to a specific protein pocket from a large database of molecules. Traditional docking methods are time-consuming, while learning-based approaches supervised by high-precision conformational or affinity labels are limited by the scarcity of training data. Recently, a paradigm of feature alignment through contrastive learning has gained widespread attention. This method does not require explicit binding affinity scores, but it suffers from the issue of overly simplistic construction of negative samples, which limits their generalization to more difficult test cases. In this paper, we propose Drug-TTA, which leverages a large number of self-supervised auxiliary tasks to adapt the model to each test instance. Specifically, we incorporate the auxiliary tasks into both the training and the inference process via meta-learning to improve the performance of the primary task of virtual screening. Additionally, we design a multi-scale feature based Auxiliary Loss Balance Module (ALBM) to balance the auxiliary tasks to improve their efficiency. Extensive experiments demonstrate that Drug-TTA achieves state-of-the-art (SOTA) performance in all five virtual screening tasks under a zero-shot setting, showing an average improvement of 9.86\% in AUROC metric compared to the baseline without test-time adaptation.

Lay Summary:

How can we make AI drug discovery models work reliably on proteins and molecules they’ve never seen before? This is a major challenge in structure-based virtual screening, where most models struggle to generalize beyond training data.

We tackle this problem by introducing Drug-TTA, a method that allows the model to adapt at test time—without access to labels. Our key idea is to leverage a set of self-supervised tasks to fine-tune the model on each test sample individually. Additionally, we use meta-learning during training to ensure that the model can quickly adapt at inference time.

Surprisingly, this simple approach significantly boosts performance across diverse benchmarks. Our results show that even without test-time labels, a model can still learn to “specialize” to new proteins and molecules. Drug-TTA thus offers a practical and data-efficient way to improve virtual screening in real-world drug discovery.

Primary Area: Applications->Health / Medicine

Keywords: Drug Virtual Screening; Test-Time Adaptation; AI for science

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Link To Code: https://github.com/ShenAoAO/Drug-TTA.git

Submission Number: 3476