Drug-TTA: Test-Time Adaptation for Drug Virtual Screening via Multi-task Meta-Auxiliary Learning
Authors: Ao Shen
Mingzhi Yuan
Yingfan MA
Jie Du
Qiao Huang
Manning Wang
Venue: ICML 2025 poster (2025)
Keywords: Drug Virtual Screening; Test-Time Adaptation; AI for science
Primary area: applications->health_medicine

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.

Drug-TTA:
Test- Time Adaptation for Drug Virtual Screening via Multi-task Meta-Auxiliary Learning
Ao Shen 12
Mingzhi Yuan 12
Yingfan Ma
Jie Du 12
Qiao Huang Manning Wang'
Abstract
1. Introduction
Virtual screening is a critical step in drug dis- covery, aiming at identifying potential drugs that bind specific protein pocket from large database of molecules. Traditional docking meth- ods are time-consuming, while learning-basse baannd ap- proaches supervised by high-precision conforma- tional or affinity labels are limited by the scarcity of training data. Recently, paradigm of fea- ture 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 con- struction of negative samples, which limits their generalization to more difficult test cases. In this paper, we propose Drug-TTi TTA, which lever- ages 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. Addi- tionally, we design a multi-scale feature based Auxiliary Loss Balance Module (ALBM) to bal- ance the auxiliary tasks to improve their effi- ciency. Extensive experiments demonstrate that Drug-TTA TTA achieves state-of-the-aus (SOTA) per- formance in all five virtual screening tasks under zero-shot setting, showing an average improve- ment of 9.86% in AUROC metric compared to the baseline without test-time adaptation. The code is available at https: : / ht////ttt/b////////// com/ ShenAoAO/ AeeenAAAO/DDu-TTTe /er///ruuuTaaa/ TTA.
Virtual screening plays crucial role in early drug discov- ery by enabling the rapid identification of potential drug candidates from vast small molecule libraries for further validation (Patel et al., 2021: Schneider, 2010). The num- ber of molecules in these libraries is increasing very fast, with recent ones reaching billions of molecules (Zhou et al., 2024). As a result, there an urgent need for fast and efficient virtual screening methods to keep pace with the growing size of molecule libraries (Sadybekov & Katritch, 2023; Shen et al., 2024a).
Traditional virtual screening methods (Halgren et al., 2004; Trott & Olson, 2010; Spitzer & Jain, 2012; Combs et al., 2013) primarily rely on molecular docking software to pre- diet binding conformations and estimate binding affinity by calculating the binding energy. However, these methods are often computationally expensive and time- coneuccnn as they require sampling large number of conformations and evaluating each through complex scoring calculations (Kitchen et al., 2004). With the advancement of deep learn- Yuan et al., 2023; Shen et al., 2024c), the majority of current methods focus on using deep models to predict dock- ing conformations (Cai et al., 2024: Zhang et al., 2023b) and binding affinity (Zhang et al., 2023a; Kimber al., 2021), which significantly accelerates the screening process. How- ever, these models are typically trained on positive pairs, i.e., binding protein-molecule pairs, with available confor- mational or binding affinity labels (Wang et al., 2005). As a result, the scarcity of high-quality labeled data limits their generalization ability. Furthermore, the absence of negative pairs, i.e., non-binding protein-molecule pairs, makes the training data have large disparity with the inference data, which requires the model to identify active molecules that form positive pairs with the target protein from a large num- ber of inactive molecules. Recently, a new mmucture-based feature alignment paradigm (e.g., DrugCLIP (Gao et al., 2024)) of virtual screening has emerged, which is free of complex docking simulation and high-quality affinity la- bels. Similar other feature alignment approaches (e.g.. CLIP (Hafner et al., 2021)), this paradigm maps molecules and pockets into a shared feature space using deep neu- ral networks and employs contrastive learning for training,
ing(Yuan
"quual contribution 'Digital Medical Rescarch Center, School of Basic Medical Sciences, Fudan University. 131 Dong' an Road, 200032, Shanghai, China 'Shanghai Key Laboratory of Medical Image Computing and Computer Assisted Intervention, Fudan University, 131 Dong's an Road, 200032, Shanghai, China. Corre- spondence to: Manning Wang <mnwang @ fudan. edu.....
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Sereening
istic. Given that drug screening often presents a zero-shot problem, where testing involves no labeled target-domain data, tuning training data to simulate the target domain is also not feasible. As Vladimir Vapnik (Vapnik, 2006) fa- mously stated, "When solving a problem of interest, do not solve a more general problem as an intermediate step. Try to get the answer that you really need, not a more gen- eral one." An alternative approach is to design a network that adapts to each test sample individually. We draw in- spiration from the success of test-time adaptation (TTA) in image classification tasks (Sun et al., 2020; Xiao & Snoek, 2024), which adapts a trained model to each test instance during inference to make better predictions. TTA typically leverages superviee auxiliary tasks during inference to dynamically update the model parameters, enabling the learning of feature representations tailored individual test instances. This approach requires prior knowledge of the test data distribution and instead adjusts model parameters dynamically on a per-instance basis during inference. In this paper, we propose Drug-TTA, which incorporates the TTA approach into DrugCLIP (Gao et al., 2024) and sig- nificantly boosts its performance on drug virtual screening. This is the first work introduce TTA into drug screening to adapt the trained model to unseen instances at test time, and the significant performance improvement demonstrates its promise.
36
Aooaage Fomure Ditiau butwose PaskieaNegeeo ged pals
(a) Training
(b) Inference
Figure/ 1. Average feature distance between positive/negative pairs the training and inference stages. Positive and negative pairs re- fer to binding and non-binding pair of protein pocket and molecule, respectively. and o points represent the average feature distance for positive and negative pairs, respectively. (a) Training stage (using training dataset in DrugCLIP (Gao et al., 2024)): Each batch contains 48 positive pairs and 2,256 (48 x 47) negative pairs. (b) Inference stage (using DUD-E benchmark (Mysinger et al., 2012)): There 102 packets and for each pocket number of molecules are provided to form positive and negative pairs with it.
where the binding pocket-ligand data in a batch are treated as positive pairs and other pocket-ligand combinations in the same batch are regarded as negative pairs. During in- ference, it scores and ranks candidate molecules based on the similarity between pocket and molecule features and se- lects the molecules with high scores as the screening results. However, this method explores the CLIP-based approach to construct the positive and negative pairs for model training, which leads to an inherent limitation in drug virtual screen- ing, Specifically, negative pairs for a given pocket are defined too simplistically, consisting of only molecules binding to other pockets. In actual virtual screening tasks, there are often challenging negative pairs, where the bind- ing and non-binding molecules for a specific pocket exhibit very similar physicochemical properties. This results in domain gap between the training and inference stages, mak- ing it challenging for the model to generalize effectively. To illustrate this problem, we conduct simple statistical anal- ysis for the training dataset in DrugCLIP (Gao et al., 2024) and the DUD-E benchmark (Mysinger et al., 2012). We calculate the average feature distance between positive pairs and negative pairs of each training batch, and the results are shown in Figure I(a). The results reveal that distinguish- ing positive and negative pairs during training is relatively easy. Similarly, we calculate the average feature distance between positive pairs and negative pairs for 102 pockets in the DUD-E benchmark dataset, and the results are shown in Figure 1(b). We can see that it is much more difficult to distinguish positive and negative pairs since their feature distances are much closer than that in the training stage.
To leverage self-supervised learning- based TTA techniques in virtual screening, we introduce five self-supervised auxil- iary tasks for proteins and small molecules, encompassing both fine-grained and pears--grained tasks. To prevent the model from overly favoring the auxiliary tasks, which could negatively impact the performance of the primary task and result in suboptimal outcomes, we explore a meta-learning framework in Drug-TTA. During training, this framework simulates the inference process by first updating the model weights using auxiliary tasks and then applying the updated model to the primary task. This process is consistently followed throughout both training and inference, ensuring that the primary task's performance is effectively optimized by the auxiliary tasks, while mitigating overfitting to the auxiliary tasks. Furthermore, to better balance the auxil- iary tasks, we design a multi-scale feature based Auxiliary Loss Balance Module (ALBM), which flexibly adjusts the weights of the auxiliary task loss during inference. This method offers greater flexibility than manually designing loss weights, and it is more effective than learning fixed weights and can generalize more effectively to new input samples during inference.
Our main contributions are summarized as follows:
To address this issue, one naive approach could be to en- hance the training dataset with more challenging negative pairs. However, designing negative pairs for training that are consistent with those encountered at test time unreal-
To the best of our knowledge, Drug-TTA is the first to introduce TTA into the field of drug virtual screening so that the trained model can be dynamically adapted to each test instance during inference.
2

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
We introduce five self-supervised auxiliary tasks for TTA in drug virtual screening, spanning both fine- grained and coarse-grained tasks, enabling the model to effectively adapt to test instances. Moreover, we propose a meta-auxiliary learning strategy that allows the model to learn how to leverage these auxiliary tasks for parameter adjustment, thereby improving the per- formance of the primary task during inference.
form negative pairs with the given pocket. A recent feature alignment paradigm (Gao et al., 2024) avoids the need for ground-truth affinity value by projecting the pocket and the molecule representation into a common feature space and scoring molecules directly based on their feature similarity to the target pocket. However, this approach suffers from overly simplistic negative pairs, which is inconsistent with the real-world virtual screening task. To solve this issue, we propose self--sue learningbbased TTA approach to dynamically adjust the model to each test instance, aiming to enhance the performance in real virtual screening tasks.
We propose multi-scale feature based ALBM to balance the weights of self-suuperviedd auxiliary task losses according to the test instance features, which outperforms directly assigning fixed weights.
Test-Time Adaptation: TTA is an emerging learning paradigm that adapts the model to target data at test time (Xiao & Snoek, 2024). TTA encompasses a variety of ap- proaches (Xiao & Snoek, 2024), and the most widely used approach is based on auxiliary tasks, which have found ex- tensive application in fields such as image processing (Yeo et al., 2023; Bahmani et al.; Park et al., 2024), video analysis (Lin et al., 2023; Xiong et al., 2024; Yi et al., 2023), 3D classification (Shim et al., 2025; Wang et al., 2024; Shin et al., 2022) and so on. Most studies utilizing this approach focus on how to design self-supervised learning tasks as auxiliary tasks and how to couple them with the primary task. For auxiliary task design, Varsavsky et al. (Varsavsky et al., 2020) introduce a framework combining adversarial loss with consistency regularization to enhance model adap- tation at test time. CKEPE (Liu et al.) adapts the model by utilizing multimodal representation learning, and NC-TTT (Osowiechi al., 2024) enhances seff-adaptation by apply- ing noise discrimination to the features. Besides designing proper auxiliary tasks, another important issue is how to couple them with the primary task to avoid overfitting the auxiliary tasks while harming the performance of the pri- mary task. For example, Liu et al. (Liu et al., 2023) propose a novel meta-auxiliary learning framework aimed at opti- mizing model adaptation specifically for test-time scenarios. TT++ (Liu et al., 2021) proposes an online feature align- ment strategy to align feature distributions at test time with those at training time, effectively mitigating overfitting to the auxiliary task. Point-TTA (Hatem et al., 2023) trains a meta-auxiliary learning TTA approach for point cloud reg- istration. In this paper, we propose an adaptive algorithm for balaneing self-supervised task losses, and utilize a meta- learning training framework to avoid bias towards auxiliary tasks.
Drug-TTA TTA achieves SOTA performance with signifi- cant improvements across five virtual screening tasks, showcasing its remarkable potential in drug virtual screening, particularly with an average improvement of 9.86% in the AUROC metric.
2. Related Work
Drug Virtual Screening: The goal of virtual screening is to identify the most relevant molecules from a large molecu- lar library that have the highest probability of binding to a given protein pocket. Drug virtual screening methods can be broadly categorized into traditional methods and deep learning-baseee based methods (Oliveira et al., 2023). Traditional methods (Halgren et al., 2004; Trott& & Olson, 2010; Spitzer & Jain, 2012: Combs et al., 2013) rely on docking soft- ware to assess the binding affinity between proteins and ligands by extensively sampling and evaluating docking conformations. However, these methods are heavily de- pendent on the accuracy of scoring functions and require complex structural sampling and energy calculations, which reduce computational efficiency. In contrast, deep learning- based methods can significantly accelerate the screening process (Shen et al., 2024b; Du et al., 2025). Most of these methods train a deep model to predict the binding affin- ity between a molecule and a pocket in a supervised way with available binding affinity labels (Öztürk et al., 2018: 8: Zheng et al., 2019; Jones al., 2021) or calculating bind- ing scores from known complex conformations (Cai et al., 2024: Zhang et al., 2023b). Then the trained model can be used to predict the binding affinity between a target pocket and molecules in a large library and rank the molecules to determine the best candidates. There are two issues in this kind of method. First, it is usually difficult obtain the ground-truth affinity labels or conformation for supervised learning. Another less obvious issue is that they rely positive pocket-ligand pairs for model training, and lack the learning of negative pairs. This makes them less aligned with most real virtual screening tasks, where the objec- tive is to select molecules that can bind to a given pocket from large pool of candidate molecules, most of which
3. Method
Given protein pocket p and set of candidate molecules M = (m1, min.,,..... m.)) ). the goal of virtual screening is to rank M according to their probability of binding with p. In this work, we propose Drug- TTA, a test-time adaptation (TTA)-based virtual screening framework, which consists of a pocket encoder and molecule encoder, parameterized
3

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Sereening
Pipeline
b. Auxiliary Branch
Fite-gralted tasks
Recomstattion
Auxillary Branch
Encoder
Takan hype
Primary Branch
Coardiatte
Encader
HAA0
Σ Caur
Pocket
Coome-grained Taaks
RL dhergercc
Excoder
Molecute
ppt
CNT
Audilarry Loss Balance Module
Auxiliary Branch
Rachword
HLP
Task specifi
Wetgits
C. Meta-learning TTA process
Training Stage
St Audiaory Task
Stop2: Updating Layer Nom
Step3: Primery Task
Stap4: Updating Encoder
Pocket
Encoder
Auciiory Branch
PP
Pocket
Eyper Nome
Sitees
"Asciliory
Pocket
Updated Encoder
Puxiiary Branch
(YRKP))
Cyrt
Pocket
Encoder Updoting
Austliary
Branch
Primary
Primary
Primary Branch
Primery
....epr
Branch
Updaring
Moleculs
Encoder
Auctiory
Branch
Layer Nera
Others
Branch
Matecula
Updated
Encoder
Ausiilory
scome
Encoder Updating
Austliary
Testing Stage
RECOAD
The framework comprises two branches: primary branch and an auxiliary branch. The input pockets and molecules are first encoded by their respective encoders. Then, the primary branch computes similarity score between the encoded features of the pockets and the molecules to train the encoders through contrastive loss. (b) Auxiliary branch: Pockets and molecules go through their own auxiliary branches, which have identical architecture but different parameters and inputs. Both auxiliary branches perform five self- supervised tasks, consisting of three fine-graiee fine-grained tasks and two goarse-grained tasks. Seven losses are calculated from the five auxiliary tasks and they are weighted and balanced through the multi-scale feature based Auxiliary Loss Balance Module (ALBM). This module concatenates features from the encoder's 's top, middle, and final layers, reduces their dimensionality through an MLP, and employs task-specific MLPs to generate adaptive weights for each loss. (c) Meta-learning TTA process: The training stage consists of four steps: (I) The pocket-ligand pairs are encoded through their respective encoders, and auxiliary task loss is computed. (2) The auxiliary loss is used to update the encoder's layer norm weights. (3) The updated encoder is employed to perform the primary task and calculate the primary loss. (4) All encoders are updated again by the primary loss. During the testing stage, the same pipeline is followed, and for each test instance, the encoder weights are first updated by the auxiliary loss, and then the final similarity score is calculated by the primary branch, which is used rank the molecules.
Figure 2. Overview of Drug- TTA. (a) Pipeline:
by 8 and 9m respectively. The two encoders are trained to extract pocket and molecule features: For (p) - fp: Foo (m) + fm. and the molecules are ranked by the sim- ilarity between fp and fm. Drug-TTA explores the TTA technique to enhance the model's generalization ability and overall performance. The model architecture and the train- ing and testing workflow are illustrated in Figure 2, and it is further explained the following three subsections.
learning paradigm is briefly introduced below for complete- ness.
First, both protein pockets and molecules are represented at the atomic level, with atom types and 3D coordinates fed into the protein/molecule encoder. In this study, we follow the encoder architecture of Uni-Mol (Zhou et al., 2023). Through the encoder, both atom features and pair- wise features are generated. Specifically, the atom features for pockets and molecules are encoded as fp and fm. re- spectively. These representations are subsequently passed through simple non-linear projection heads ofre and ore- yielding the final output features f and fr.
In Section 3. we introduce the primary branch (Figure 2(a)), which learns feature alignment between pocket and molecule features through contrastive learning. Section 3.2 introduces the five auxiliary tasks and explains how we perform weight learning to balance them (Figure 2(b)). Finally, in Section 3.3, we describe the training and testing TTA process (Figure 2(c)).
During training, a contrastive learning approach is ex- plored, where positive pairs are the input binding protein- molecule pairs. Similar to CLIP (Hafner et al., 2021), an in-batch sampling strategy is employed to generate negative pairs. Specifically, given a batch of N input positive pairs {(=2 xx)) > with the list of pockets { and the list of molecules { The pockets and molecules can be combined to N2 pairs (xf, 2)"), where i.j e [1.
3.1. Primary Branch
In this study, we follow the contrastive learning paradigm of DrugCLIP (Gao et al., 2024) to construct the virtual screening primary branch to perform the primary task. The
4

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Sereening
and pairs with i are positive pairs, while those with i # serve as negative pairs.
balance the distribution of weights, we apply a softmax function constrain the weights, ensuring their sum equals one. As a result, final weights { (1 2, .., 7) for the auxiliary task losses are mapped as follows:
The training loss consists of two parts: the Pocket-to-Met expla(uf. 2/1) and the Mol-to-Pocket
loss C = N E, exp(a(z), a(")/7)
1 exple(xf, x //7
f.
1 267
loss Ch = N E. cap(a(a) =f)/r)> where = f... is the similarity functions composed of both the dot product(-) and cosine similarity, and represents the temperature pa- rameter controling the softmax distribution. The former measures the likelihood of ranking the binding molecule higher than other molecules for a given protein, while the lat- ter evaluates the probability of correctly ranking the binding targets for a given molecule. By combining these two parts, the primary branch training loss obtained as follows:
\ Soofmmx
(3)
The total auxiliary loss is formulated as:
Lauz = =AiCrupe + >2eoooord + daadiet
(4)
+MCNX + AGCNP+ + deCm + d Ay
3.3. Meta-learning TTA
N
Our goal is to leverage self-supervised auxiliary tasks to adjust the encoder weights during inference, enabling the model to quickly adapt its parameters for each test instance without the need for additional labels. To prevent the param- eters updated by the auxiliary loss from becoming overly focused on improving the auxiliary tasks at the expense of the primary task, as suggested in (Hatem et al., 2023), we propose a meta-auxiliary task approach that jointly trains both the main and auxiliary tasks. There are three groups of parameters in the whole model, the pocket/molecule encoder parameters PP /am, the pocket/molecule auxiliary branch head parameters o / aux. and the pocket/molecule pri- mary branch output head parameters o / pri We define three sets of parameters according to their different updating mechanism in training and inference. The updated param- eter set in primary branch is pri = {0P, 0"" ofprs o m while the updated parameter sets in protein and molecule auxiliary branch are @ Pan. fohormy oll omurr } and fm...
Cpri
(CE + Cg).
(L)
3.2. Auxiliary Branch
In this work, we introduce five auxiliary self- supervuee tasks in the auxiliary branch adapt the encoders differ- ent data during inference. These tasks consist of three fine- grained tasks focusing on contextual information and two coarse- goarseegriindd tasks focusing on modeling global features. The fine-grained tasks are masked atom type prediction, corrupted coordinate reconstruction and distance prediction, and the pparse-grained tasks are KL, divergence minimiza- tion and SimCLR contrastive learning (Chen et al., 2020). These tasks generate seven distinet losses, with detailed descriptions provided in the Appendix A.
To better balance the aforementioned self-supervised tasks and their losses, one naive approach is to learn the weights for different losses during training and fix them for testing. However, this leads to overfitting of the weights to the train- ing data, which contradicts the purpose of TTA-to adjust the model for better alignment with testing conditions. To this end, we propose the multi-scale feature based ALBM, which adjusts the weight allocation at test time based on the test instance's 's own feature.
{ (othh norm: ow Face) J. where Pnorm and a norm are the parame- ters of normalization layers in l and 6m..
As shown in Figure 2(c), the training and testing processes are largely similar, with the key difference occurring after the primary task is executed. During training, the encoder weights are updated based on the primary task loss by ad- justing épri. In contrast. during testing, the model directly outputs the similarity score for the primary task without further updating the encoder weights.
As shown in Figure 2(b), a transformer architecture is used for both the protein and the molecule encoders, which con- sist of multiple layers. To capture comprehensive repre- sentation of test instance, we extract features from the top, middle and last layers of the encoder, which are denoted as ftop: fmii and flast, respectively, and concatenate them. Subsequently, the concatenated features are passed through seven independent MLP heads [MLP;} (1 = 1, 2,,,.., 7). each responsible for calculating a weight for one of the seven losses. The weights {et} (I 1,2,..... ..., 7) are calcu- lated as:
The initial weights for both encoders are obtained from pre-training on a large-scale dataset by Uni-Mol (Zhou et al., 2023). A batch of sample pairs (< =x > in- put, and the molecules and the pockets are separated and fed into their respective encoders and auxiliary task branches. The normalization layer parameters @horm and @horm as well as head parameters 0% and o are updated based on respective auxiliary task losses 2 and Lm
Training:
e1 MLP; (Concat(frop Imidd flast))
(2)
( Th dlue x- Whorrm of duur) ) a VoL fomorm o mm x fo o a V o
(5)
To prevent the model from collapsing during training and
5

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
where a is the learning rate for auxiliary tasks.
Table I. Results on DUD-F. AUROC, BEDROC, and EF are re- ported (higher values indicate better performance). Bold values represent the best performance, and green indicates improvements of Drug-TTA over DrugCLIP.
As the primary branch has been updated by o Shorm and @ the primary task loss is used to update all parameters of the primary branch:
ppri- + 0pri - VV Cpri:
Method
AUROC(S)
BEDROC(%)
0.30%
EF
1%
5%
Gilid-SP
76.70
40.70
19.39
16.18
7.23
where 3 is the learning rate for the primary task.
Vina
71.60
9.13
7.32
4.44
NN-soore
68,30
12:20
4.16
4,02
3.12
Testing: During testing, all trained model parameters are used. When pocket p and a set of candidate molecules M are input, the normalization layer parameters as well as head parameters are first updated via the auxiliary tasks. Then, the updated encoders are directly used to perform the primary task. For the inference of a new iteration, the original parameters are reloaded to prevent accumulation offsets. Pseudo code for the training and the testing process listed in Appendix C.
RPscore
65.21
12.41
4.90
4.52
2.98
Pafmay
63.11
16.50
4.24
3,86
3.76
OnionNct
2,84
Planet
71.60
10.23
8.83
5.40
DrugCLIP
80.93
38.07
31.89
10.66
93.16
82,82
57.50
54.04
16.88
DTug-TTA
(112.23)
(1332330
(111.43)
(122.15)
(16.22)
per pocket. Each active molecule corresponds to 50 decoys, which are inactive molecules with similar physicochemical properties but different 2D topologies.
As shown in Table we evaluate our method in zero- shot setting and compare its performance with docking- based (Glide-SP (Halgren et al., 2004), Vina (Trott & Ol- son, 2010)), learning-based (NN-score (Durrant & McCam- mon, 2011), RFscore (Ballester & Mitchell, 2010), Pafnucy (Stepniewska- -zzibbis-zz et al., 2018), OnionNet (Zheng et al., 2019), Planet (Zhang et al., 2023a), and feature align- ment method (DrugCLIP (Gao et al., 2024)). The results clearly indicate that the feature alignment paradigm out- performs the other paradigms, showing that this paradigm is more suitable for virtual screening. Compared to Drug- CLIP, our method demonstrates a substantial performance improvement on all metrics, particularly in BEDROC, where we can observe an increase of 32.30%. This highlights that the TTA technique significantly enhances the model's gen- eralization performance on unseen data.
4. Experiments
To evaluate the performance of our method, we first as- sess the zero-shot performance of Drug-TTA on five virtual screening benchmarks: DUD-E (Mysinger et al., 2012), LIT-PCBA (Tran-Nguyen et al., 2020), AD (Chen et al., 2019), DEKOIS 2.0 (Bauer et al., 2013), and CASF-2016 (Su et al., 2018), comparing it with existing methods in Section 4.1. Here, "zero-shot" refers to the setting where the model is tested on entirely unseen protein pockets with- out using any labeled samples from the target benchmarks. Each virtual screening benchmark consists of multiple pock- ets, with each pocket corresponding to various active and inactive molecules (decoys). The task of virtual screening is to rank molecules such that active ones that bind to the corresponding pocket are placed at the top. To evaluate screening effectiveness, we use AUROC, BEDROC, and EF as our evaluation metrics. BEDROC is an improved version of AUROC, incorporating exponential weighting to emphasize to early rankings, while EF is widely used virtual screening metric. Detailed definitions of evaluation metrics are in Appendix F. For training, we use the dataset from DrugCLIP (Gao et al., 2024), and when evaluating on each benchmark, we exclude from the training datasets and retrain TTA. The hyperparameter settings for both training and testing are provided in Appendix E. Further- more, we conduct ablation studies on key components, detailed in Section 4.2. The visualization results can be found in Section 4.3.
LIT-PCBA Benchmark (Tran-Nguyen et al., 2020): LIT- PCBA (Large-scale Interaction and Toxicity Prediction for Chemical Bioactivity Assays) is more challenging bench- mark for virtual screening, designed to address the data bias issues found in other benchmarks such as DUD-E. It con- tains 15 pocket targets, 7,844 active molecules and 407,381 unique inactive molecules selected from high-confidence PubChem Bioassay data.
As shown in Table 2, we compare our method with docking- based methods (Surflex (Spitzer & Jain, 2012), GLide-SP (Halgren et al., 2004)), learning-based methods (Planet (Zhang et al., 2023a), Gnina (McNutt et al., 2021), Deep- DTA (Öztürk et al., 2018), BigBind (Brocidiacono et al., 2023)), SPRINT (McNutt et al., 2024)) and feature align- ment method (DrugCLIP (Gao et al., 2024)). The results demonstrate that our method significantly improves the per- formance of DrugCLIP by a large margin, achieving SOTA performance across all metrics, with a particularly notable improvement in the EF metric with an order-of of magnitude enhancement.
4.1. Evaluation on virtual screening benchmarks
DUD-E Benchmark (Mysinger et al., 2012): DUD-E (Directory of Useful Decoys, Enhanced) Benchmark is a widely used virtual screening performance evaluation dataset. DUD-E contains 102 protein pockets and 22,886 active molecules, with an average of 224 active molecules
6

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Sereening
Table 2. Results on L.T-PCBA. AUROC, BEDROC, and EF are reported (higher values indicate better performance). Bold values represent the best performance, and green indicates improvements of DTA over DrugCLIP.
Table 4. Ablation Study on Various Aspects. "w/o" indicates the removal of the respective component compared to Drug- -TA.
EF
Ablation espeets
Method
AUROCYS)
BEDROC/S)
8.0%
1%
5%
Method
BEDROC(%)
EF
DrgCCLPP
80:83
0.50%
1%
5%
DTg
93.66
82.82
57.50
56.04
16,88
Surflex
81.47
2.50
SP
53.15
4,00
3.17
3.41
2.01
wiighe:
81.63
51.75
78.15
55.05
11.29
Planet
57.31
4.64
3.87
2.43
Wajahn Allocaiion
foxd weight
21.33
64.00
51.69
42.00 40.40
13.09
Gnina
60.93
5.40
4.63
weight regulartzation
DeepLITA
56.27
2,53
1.47
BigBind
60,80
Self-S Self-Sypenfiod Tas Sclection
wio coerse grainsd
82.33
53.35
SPRINT
73.4
12.3
15,90
10.78
5.92
wh fine gradigggdddd tasks
79:35
41.77
31.58
26.05
9.53
DragCLIP
57.17
6.23
8,56
2.27
Ausiliar Branch Selection
with predez TTT
81.64
55.00
40.41
35, 35
11:90
71.24
45.08
74.39
42.74
10.61
DTag-TTA
(114.07)
(138.85)
(165.83)
(137.23)
(18.34)
with madeoatl TTA
79:85
41.55
11:34
76, 06
middle and list
83.55
50.12
43.54
36,95
12.47
wio top and layers
84.35
59.78
36.31
12.56
Matiiscale Feoture Selection
wio top and riddile layes
$4.19
20.67
44.25
38.67
12.51
Table Results on AD, DEKOIS 2.0, and CASF-2016. AUROC, BEDROC, and EF are reported (higher values indicate better per- formance). Bold values represent the best performance, and green indicates improvements of TTA over DruPCLIP.
layer
9962
68.64
48.35
45.89
14.59
sop liyer
who riiitta layer
$1.60
62.00
90.94
43.92
13:47
Meu-lcuning
wil mott fourning
80.19
01.00
31.19
25.14
9.46
Dutaset
Method
AUROC(S)
BEDROCYS)
0.50%
DesgCLIP
81.19
52.04
20.50
18.00
9.10
molecules and 280 inactive molecules. Table 3 shows that Drug-TTA outperforms DrugCLIP across all metrics, demonstrating that TTA effectively enables the model to
AD
Drug TTA
(111.43
92.02
(T34,60
86.73
(112.04)
32.54
(12.51)
36.51
(76.19)
15:29
DesgCLIP
77.98
47.32
18:48
17.02
8.52
DEKOIS 20
TTA
(15.6)
M3.61
(126.30)
73.64
(7795)
26.41
(18.52)
25.64
(1454)
13.06
adapt to unseen datasets.
DrugCtl IP
85.90
67.88
36.19
$4.19
12:93
CASE 2016
Drup: TTA
(15.96)
91.88
(117.84)
85,73
('6.4)
42.75
(1701)
41.34
16.04
4.2. Ablation study
We conduct ablation experiments on the allocation of weights across different auxiliary task losses, and the results are shown in Table 4. First, we evaluate the model's performance without any weight allo- cation for the self-supervised tasks by setting all weights to (the row "weights=1"), and the results indicate that the improvement over the baseline model (DrugCLIP) is minimal. Next, we evaluate the model's 's performance when using fixed weights after training, as shown in the row "fixed weight". Compared to DrugCLIP, we observe significant im- provements across most metrics, except for AUROC, though the improvements are not as substantial as those seen with Drug-T TTA. Furthermore, we assess whether using regular- ization constraints the weights affects performance. The results in the row "w/o weight regularization" demonstrate that the absence of regularization leads to the failure of the TTA strategy, highlighting the importance of weight regularization.
Weights Allocation:
AD Benchmark (Chen et al., 2019): This dataset is an improvement over DUD-E. Specifically, it reduces negative target bias by using active molccules from other targets as decoys for the current target. For each target, active molecules from 101 other targets are selected and docked in 101 separate batches, with the top 50 molecules by affinity chosen in each batch. This results in average of 5,000 molecules used as decoys for each target. The comparison of our method with DrugCLIP is shown in Table 3. It can be observed that applying the TTA strategy significantly enhances the model's performance. When comparing the results in Table 1, we find that the improvements in the AD dataset over DUD-E have minimal impact on the AUROC and BEDROC metrics for both DrugCLIP and Drug-TTA, but they significantly affect the EF metric.
DEKOIS 2.0 Benchmark (Bauer et al., 2013): DEKOIS 2.0 is a dataset specifically designed for evaluating virtual screening methods in drug discovery. consists of tar- gets from different protein families, with each target having 40 active ligands and 1 1,200 decoys. As shown in Table 3, Drug-TTA outperforms DrugCLIP across all metrics, with particularly notable improvement of 26.32% in BEDROC.
Self- Self-Supervised Task Selection: Given the large number of tasks used in Drug-TTA, we focus on ablation experi- ments involving the coarse-grained and fine-grained tasks. As shown in Table 4, using only fine- ane-ggrined tasks ("w/o coarse-grained tasks") still results in performance improve- ments over DrugCLIP, although the gains are smaller. In contrast, when only coarse- coarse-grained tasks are used, the per- formance decreases compared to DrugCLIP. Significant per- formance gains are only achieved when both coarse- -rained
CASF-2016 Benchmark (Su et al., 2018): CASF-2016 Benchmark provides 57 targets for virtual screening tasks, with each target containing approximately active
7

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
DrugCLIP
Drug-TTA
DrugCLIP
Drug-TTA
* Pockt
Negative mollccte
AUROC-ST 40% BEDROC
AUROCCRR 12% REDROC 54%
ALIROC-6 99% BEDROCOOIN
AUROC:91 49%
BEDROC-AS.54 545.5
(a) CP3A4
(b) DHII
Figure 3. Visualization of feature distributions for the targets CP3A4 and DHT1 (b) using DrugCI.PP and Drug-TTA models. The red points represent pocket features, blue points denote active molecule features, and gray points indicate inactive molecule features.
Sumple
Mideeak veBdi
Log!
Riratak beads
Cariee
greater importance than those from the top and last layers. For further insights of these three layers, additional t-SNE visualizations refer to Appendix G.
abre CAER
8145
01463 E2154
31215
29:34
4.2
15
25
To verify that meta-learning can effectively prevent the model from becoming biased toward the auxil- iary task and producing suboptimal results, conduct an ablation experiment on the meta-learning approach, with the results shown in Table 4. The results indicate that without the meta-learning based TTA approach (w/o meta-learning), the model's ' performance is even lower than DrugCLIP, sug- gesting that this approach leads to performance decline in the primary task.
Meta-learning:
401,7
17
Figure 4. Weight allocation heatmap of four molecules TTA and their molecular properties. The heatmap illustrates the weight allocation for four molecules during the TTA process. The cor- responding molecular properties include molecular weight, LogP, the number rotatable bonds, and carbon atom count, shown alongside the molecular structures.
4.3. Visualization
We visualize the embeddings of two pocket targets and their corresponding molecules from the DUD-E dataset in Fig- ure 3. The t-SNE plot illustrates both the pocket features and their corresponding active and inactive features. The visualization clearly demonstrates that, compared to Drug- CLIP, Drug-TTA active and inactive molecule features. Notably, the active molecule points are more tightly clustered around the pocket feature points, indicating that Drug-TTA improves feature alignment. Moreover, the performance metrics, including AUROC and BEDROC, show substantial improvements, with BEDROC increasing by order of magnitude. This effect aligns with our objective of using feature alignment to identify molecules similar to given pocket feature for virtual screening.
and fine-grained tasks are combined (Dug-TTAA, demon- strating the value of combining these task types for a more comprehensive learning of instance features. More details on the combination ablation of self-supervised tasks can be found in Appendix D)
TTA significantly enhances the separability of
Drug-TTh TTA applies TTA to fine-tune the encoders for both the pocket and molecule branches. As shown in Table 4, we evaluate the model's per- formance when TTA is applied to only one encoder-ciile the molecule encoder ("wwo pocket TTA") or the pocket encoder ("w/o molecule TTA"). The results show that cluding TTA from either encoder results in a performance decrease. Notably, applying TTA only the pocket encoder results in worse performance than DrIg indicating that model adjustment requires simultaneous fine-tuning of paired data encoders to achieve optimal results.
Auxiliary Branch Selection:
Furthermore, we analyze the weight allocation across dif- ferent molecules. Figure 4 presents the heatmap of the weight distributions for four molecules with distinct chem- ical properties. This heatmap illustrates how the weight distribution varies according to different tasks, with the model adjusting the weight based on the inherent features of the test molccules. The varying weight assignments across molecules highlight the model's ability to dynamically al- locate weight according to each sample, thereby allowing for more accurate focus the most relevant features to improve performance.
In the process of generating loss weights controlled by instance features, our approach concatenates features from the encoder's top, mid- dle, and last layers. We perform an ablation study the selection of these layers, and the results are shown in Table 4. The results indicate that combining features from all three layers provides the best representation of the sample's 's features. Notably, the middle layer features demonstrate
Multi-scale Feature Layer Selection:
8

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
5. Conclusion
F. M. Evaluation and optimization of virtual screening workflows with dekois 2.0-a public library of challenging docking benchmark sets. Journal of chemical information and modeling. 53(6): (166:1447-14622 1462, 2013.
In this paper, we introduce TTA for virtual screening for the first time. This approach effectively mitigates the is- sue of overly simplistic negative samples in the baseline model during training and better aligns with real-world vir- tual screening scenarios. Specifically, in such scenarios, both pockets and candidate molecules are unseen, and a large number of molecules fail to bind to the given pocket. Ablation study reveals that simply applying TTA to virtual screening is not sufficient. our proposed framework, we incorporate abundant self-supervised auxiliary tasks and de- sign an innovative module, multi-scale feature based ALBM, to generate weights for auxiliary task losses based the features of each instance at test time, effectively addressing the task balancing issue. To further leverage the informa- tion from auxiliary tasks to improve the primary task, we adopt a meta-learning strategy to better couple the training and inference processes of both the primary task and the auxiliary tasks, preventing model bias towards the auxiliary tasks, which may negatively affect the primary task's perfor- mance. Significant performance improvement over current SOTA methods is achieved by the proposed method, which shows its great potential in virtual screening applications. In addition, this study shows the importance of TTA in virtual drug screening and this finding may spur many future work in this direction.
Brocidiacono, M., Francoeur, P., Aggarwal, R., Popov, K. I., Koes, D. R., and Tropsha, A. Bigbind: learning from nonstructural data for structure- based virtual screening. Journal of Chemical Information and Modeling, 64(7): 2488-2495, 2023.
Cai, H., Shen, C., Jian, T., Zhang, X., Chen, T., Han, X., Yang, Z., Dang, W., Hsieh, C.-Y., Kang, Y., et al. Car- sidock: deep learning paradigm for accurate protein- ligand docking and screening based on large-scale pre- training. Chemical Science, (544):14491444 2024.
2024.
Chen, L., Cruz, A., Ramsey, S., Dickson, C. J., Duca, J. S., Hornak, V., Koes, D. R., and Kurtzman, T. Hidden bias in the dud-e dataset leads to misleading performance of deep learning in structure-based virtual screening. PloS one, 14(8)xc0220113, 2019.
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual rep- resentations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Combs, A., Deauca, S. L.,, DeLuca, S. H., Lemmon, G. H., Nannemann, D. P., Nguyen, E. D.,, Willis, J. R., Sheehan, J. H., and Meiler, J. Small-molecule ligand docking into comparative models with rosetta. Nature protocols, 8(7):1277-1298. 2013.
Acknowledgments
This work is supported by the Science and Technology In- novation Plan Of Shanghai Science and Technology Com- mission grant RS$41900400, Fudan University Science In- telligence Special Fund grand FD-AI4S04183.
Du, J., Yuan, M., Shen. A., and Wang, M. Ppdock: Pocket prediction- protein- blind docking. Journal of Chemical Information and Modeling, 2025.
Impact Statement
Durrant, J. D. and McCammon, J. A. Nnscore 2.0: a neural-
network receptor receptor-ligand scoring function. Journal of chemical information and modeling, 51811)38877-2903 2011.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Gao, B., Qiang, B., Tan, H., Jia, Y., Ren, M., Lu, M., Liu, J., Ma, W.-Y., and Lan, Y. Drugelip: Contrasive protein- molecule representation learning for virtual screening. Advances in Neural Information Processing Systems, 36, 2024.
References
Bahmani, S., Hahn, O., Zamfir, E., Araslanov, N., Cremers, D.. and Roth, S. Semantic self-adaptation: Enhancing generalization with single sample. Transactions Machine Learning Research.
Hafner, M., Katsantoni, M., Köster, T., Marks, J., Mukher- jee, J.. Staiger, D., Ule, J., and Zavolan, M. Clip and com- plementary methods. Nature Reviews Methods Primers,
Ballester, P. J. and Mitchell, J. B. machine learning approach to predicting potein-ligand binding affinity with applications to molecular docking. Bioinformatics, 26(9):1169-11) 2010.
Halgren, T. A., Murphy, R. B., Friesner, R. A., Beard, H. S., Frye, L. L., Pollard, W. T., and Banks, J. L. Glide: a new approach for rapid, accurate docking and scoring. 2. enrichment factors in database screening. Journal of 2004.
Bauer, M. R. Ibrahim, M., Vogel, S. M., and Boeckler,
medicinal chemistry, 47())))

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
Hatem, A., Qian, Y., and Wang. Y. Point-tta: Test-time adaptation for point cloud registration using multitask meta-auxiliary learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 16494 16504, 2023.
Mysinger, M. M., Carchia, M., Irwin, J. J., and Shoichet, B. K. Directory of useful decoys, enhanced (dud-e): bet- ter ligands and decoys for better benchmarking. Journal of medicinal chemistry, 55(14):6582-6 65944, 2012.
Oliveira, T. A. d., Silva, M. P. d., Maia, H. B., Silva, A. M. d., and Taranto, A. G. Virtual screening algorithms in drug discovery: A review focused machine and deep learning methods. Drugs and Drug Candidates, 2(2): 31 311-334, 2023.
Jones, D., Kim, H., Zhang, X., Zemla, A., Stevenson, G., Bennett, W. D., Kirshner, D., Wong, S. E.,, Lightstone, F. C.,, and Allen, J. E. Improved protein-ligand binding affinity prediction with structure- based deep fusion infer- ence. Journal of chemical information and modeling, 61 (4):1583-1592. 2021.
Osowiechi, D., Hakim, G. A. V., Noori, M., Cheragha- likhani. M., Bahri, A., Yazdanpanah, M., Ben Ayed, I., and Desrosiers, C. Nc-ttt: A noise constrastive approach for test-time training. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6078-6086, 2024.
Kimber, T. B.. Chen, Y., and Volkamer, A. Deep learning in virtual screening: recent applications and developments. International journal of molecular sciences, 22(9):4435, 2021.
Özüürk, H., Özgür, A., and Ozkirimli, E. Deepdta: deep drug taaget binding affinity prediction. Bioinformatics, 34(17):1821-i82) 2018.
Kitchen, D. B., Decornez, H., Furr, J. R., and Bajorath, Docking and scoring in virtual screening for drug discovery: methods and applications. Nature reviews Drug discovery. 3(11):935-94) 2004.
Park, H., Gupta, A., and Wong, A. Test-time adaptation for depth completion. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 20519-20529, 2024.
Lin, W., Mirza, M. J., Kozinski, M., Possegger, H., Kuchne, H., and Bischof, H. Video test-time adaptation for action recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22952- 22961, 2023.
Patel, A. R., Patel, H. B., Mody, S. K., Singh, R. D., Sar- vaiya, N., Vaghela, S. H., and Tukra, S. Virtual screen- ing in drug discovery. Journal of Veterinary Pharmacol- ogy and Toxicology, 10(2):1-9, 2021.
Liu, C., Wan, Z., Ouyang, C., Shah, A., Bai, W., and Ar- cucci, R. Zero-shot ecg classification with multimodal learning and test-time clinical knowledge enhancement. In Forty-first International Conference on Machine Learn- ing.
Sadybekov, A. V. and Katritch, V. Computational ap- proaches streamlining drug discovery. Nature, 616(7958): 673-685, 2023.
Liu, H., Chi, Z., Yu, Y., Wang, Y., Chen, J., and Tang, Meta-auxiliary learning for future depth prediction in videos. In Proceedings of the IEEF/CVF Winter Confer- ence on Applications of Computer Vision, pp. 5756-5765, 2023.
Schneider, G. Virtual screening: an endless staircase? Na- lure Reviews Drag Discovery, 9(4):273-27,, 2010.
Shen, A., Yuan, M., Ma, Y., Du, J., and Wang, M. Comple- mentary multi modality molecular self-supervised Iearn- ing via non- aonovrrlapping masking for property prediction. Briefings in Bioinformaties, 25(4)bbbae256, 2024a.
Liu, Y., Kothari, P., Van Delft, B., Bellot-Gurlet, B., Mordan, T., and Alahi, A. Tit++: When does self-supervised test-time training fail or thrive? Advances Neural Information Processing Systems, 34:21808- 21820, 2021.
Shen, A., Yuan, M., Ma, Y., Du, J., and Wang, M. Pgbind: pocket-guided explicit attention learning for protein- ligand docking. Briefings in Bioinformatics, 25(5): bbae455, 2024b.
McNutt, A. T., Francoeur, P., Aggarwal, R., Masuda, T., Meli, R., Ragoza, M., Sunseri, J., and Koes, D. R. Gnina 1.0: molecular docking with deep learning. Journal of cheminformatics, 13(1):43, 2021.
Shen, A., Yuan, M., Ma. Y., and Wang, M. Ss-pro: A sim- plified siamese contrastive learning approach for protein surface representation. Frontiers of Computer Science, 18(5):185110, 2024c.
McNutt, A. T., Adduri, A. K., Ellington, C. N., Dayao, M. T., Xing, E. P., Mohimani, H., and Koes, D. R. Sprint enables interpretable and fltra-fast virtual screening against thou- sands of proteomes. arXiv preprint arXiv:241..15418, 2024.
Shim, H., Kim, C., and Yang, E. Cloudfixer: Test-time adap- tation for 3d point clouds via diffusion-guided geometric transformation. In European Conference on Computer Vision, pp. 454-471 Springer, 2025.
10

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
Shin, I,, Tsai, Y.-H., Zhuang, B., Schulter, S., Liu, B., Garg, S., Kweon, I. S., and Yoon, K.-J. Mm-tta: multi-modal test-time adaptation for 3d semantic segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16928-16937, 2022.
Xiong, B., Yang, X., Song, Y., Wang, Y., and Xu, C. codaaity-collabooative test-time adaptation for action recognition. In Proceedings of the IEEE/CVF Confer- ence on Computer Vision and Pattern Recognition, pp. 26732-26741, 2024.
Spitzer, R. and Jain, A. N. Surftex-occk: Docking bench- marks and real- world application. Journal of computer- aided molecular design, 26:687-699 2012.
Yeo, T., Kar, O. F., Sodagar, Z.,, and Zamir, A. Rapid net- work adaptation: Learning to adapt neural networks us- ing test-time feedback. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4674- 4687, 2023.
Saaniiwwska-ziibbnssee M. M.. Zielenkiewicz, P., and Siedlecki, P. Development and evaluation of a deep learn- ing model for protein- ligand binding affinity prediction. Bioinformatics, 34(21):3666-3))) 3674, 2018
Yi, C.,, Yang, S., Wang, Y., Li, H., Tan, Y.P.., and Kot, A. C. Temporal coherent test-time optimization for robust video classification. arXiv preprint arXiv:2302. 22233022.4400 14309, 2023.
Su, M., Yang, Q. Du, Y., Feng, G., Liu, Z., Li, Y., and Wang, R. Comparative assessment of scoring functions: the casf-201 update. Journal of chemical information and modeling, 59(2):895-913 2018. 8.
Yuan, M., Shen, A., Fu, K., Guan, J., Ma, Y., Qiao, Q., and Wang, M. Proteinmae: masked autoencoder for protein surface self seff-supervised learning. Bioinformaties, 39 2): btad724, 2023.
Sun, Y., Wang, X., Liu, Z.,, Miller, J., Efros, A., and Hardt, M. Test-time training with self-supervision for generaliza- tion under distribution shifts. In International conference machine learning, pp. 9229-9248. PMLR, 2020.
Zhang, X., Gao, H., Wang, H., Chen, Z., Zhang, Z., Chen, X., Li, Y., Qi, Y., and Wang, R. Planet: multi-objective graph neural network model for pigtein-ligand binding affinity prediction. Journal of Chemical Information and Modeling, 64(77)2205-222) 2023a.
Tran-Nguyen, V.-K., Jacquemard, C., and Rognan, D. Lit- pcba: an unbiased data set for machine learning and virtual screening. Journal of chemical information and modeling, 6019):4263-4273. 2020.
Zhang, X., Zhang, O., Shen, C., Qu, W., Chen, S., Cao, H., Kang, Y., Wang, Z., Wang, E., Zhang, J., et al. Efficient and accurate large library ligand docking with karmadock. Nature Computational Science, 3(9):799-80)) 804, 2023b.
Trott, O. and Olson, A. J. Autodock vina: improving the speed and accuracy of docking with a new scoring func- tion, efficient optimization, and multithreading. Journal of computational chemistry, (1(2):555-5) 4611,,466122000 2010.
Zheng, L., Fan, J., and Mu, Y. Onionnet: multiple-layer convolutional neural net- work for protein-ligant binding affinity prediction. ACS omega, 4 1554))5556-11966 2019.
peeemollcullr--conttt---baee
Vapnik, V. Estimation of dependences based on empirical data. Springer Science & Business Media, 2006.
Zhou, G., Gao, Z., Ding, Q., Zheng, H., Xu, H., Wei, Z.,, Zhang. L., and Ke, G. Uni-mol: A universal 3d molecu- lar representation learning framework. In The Eleventh International Conference on Learning Representations, 2023.
Varsavsky, T., Orbes-Arteagee M., Sudre, C. H., Graham, M. S., Nachev, P., and Cardoso, M. J. Test-time unsuper- vised domain adaptation. In Medical Image Computing and Computer Assisted Intervention-MICCA. 2020: 23rd International Conference, Lima, Peru, October 4-8, 2020, Proceedings, Part 1 23, pp. 428 436.436.775 Springer, 2020.
Zhou, G., Rusnac, D.-V., Park, H., Canzani, D., Nguyen, H. M., Stewart, L.,, Bush, M. F., Nguyen, P. T., Wulff, H., Yarov- Yarovoy, V., et al. An artificial intelligence accelerated virtual screening platform for drug discovery. Nature Communications, 15(1):7761, 2024.
Wang, R., Fang, X., Lu, Y., Yang, C.-Y., and Wang. The pdbbind database: methodologies and updates. Journal of medicinal chemistry, 48(12):4111 11111114199111 2005.
Wang, Y., Cheraghian, A., Hayder, Z., Hong, J., Ramas- inghe, S., Rahman, S., Haamedt-Arisizaate D., Li, X., Petersson, L., and Harandi, M. Backpropagation-te free net- work for 3d test-time adaptation. In Proceedings of the JEEE/CVF Conference on Computer Vision and Patlern Recognition, pp. 23231-23222 123241, 2024.
Xiao, Z. and Snoek, C. G. Beyond model adaptation at test time: survey. arXiv preprint arXiv:24/1.0300 2024.

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
A. Auxiliary losses
In this work, we introduce five auxiliary self-supervised tasks in the auxiliary branch to adapt the encoders to different data during inference. These tasks consist of three fine- geee-grried tasks focusing on contextual information and two coarse- boase-grained tasks focusing on modcling global features. The fine-gre fine-grained tasks are masked atom type prediction, corrupted coordinate reconstruction and distance prediction, and the coarse-gras gaarse-griindd tasks are KL divergence minimization and SimCLR contrastive learning (Chen et al., 2020).
Masked atom type reconstruction: This task aims to recover the atom type of masked atoms in a sample, mainly leveraging the local context provided by surrounding atoms. By doing so, it helps the encoder learn the relationships of chemical nature between atoms. The process begins by randomly masking a subset of atoms in the molecule or protein pocket. Some of the masked atoms are replaced with a special [MASK] token, while others are substituted with random tokens (Zhou et al., 2023). These masked molecules or pockets are then passed through their respective encoders to extract feature representations. Subsequently, the extracted features are fed into prediction head to predict the original atomic types at the masked positions. The loss function is based on the negative log-likelihood of the predicted atomic types at the masked positions. Formally, the loss is expressed as:
A Ltype = Σ log(p(f:(t4))
where & is the predicted atom type for the i-th masked atom, & is the ground-truth atom type, p(falt;) is the predicted probability, and A is the total number of masked atoms. Following Uni-Mol (Zhou et al., 2023), an additional loss, N is introduced to normalize masked atom features.
The following two reconstruction tasks follow this process, where the encoder output is used along with a prediction head for the reconstruction task, but with distinct output heads, each comprising a specific MLP.
Corrupted coordinate reconstruction: Building upon the previous task, this task is to reconstruct the 3D coordinates of the masked atoms, further enhancing the model's 's ability to understand spatial relationships and molecules' or pockets' geometry. Specifically, the same subset of masked atoms is considered, but their coordinates are now perturbed following the corrupted position generation and assignment algorithm of Uni-Mol (Zhou et al., 2023). The loss for this task is computed by comparing the predicted coordinates with the ground-truth coordinates and is defined as:
Leoord (SmoothLa(()cc)
(8)
where é and G are the predicted and ground-truth coordinate for the i-th masked atom, respectively.
This task aims to predict the distances between pairs of masked atoms in the pockets or molecules, thus helping the model learn how atoms interact with each other in terms of their relative spatial positioning. For this task, all pairwise distances involving masked atoms from the previous task need to be reconstructed. The loss calculated by comparing the predicted distances with the true ones and defined as:
Distance reconstruction:
Edist 1 Σ Csmothhhidd
(9)
where de and d, are the predicted and ground- ground-truth distance for the i-th pair to be reconstructed, and B is the total number of pairs that need to reconstructed. To normalize masked pairwise features, an additional loss, Lxp is introduced following Uni-Mol (Zhou al., 2023).
This task encourages the model to learn a feature distribution that closely approximates uniform distribution. As result, the model is foreed to learn more diverse feature representations, rather than excessively relying on specific patterns or features. Consequently, this regularization approach helps improve the model's 's generalization ability, enabling it to perform more effectively on unseen data. The KL loss L is computed as the KL divergence between the softmaxed logits of latent features and a uniform distribution:
KL divergence minimization:
Lu KL (ooftmar(Z)/aa || U/Ta),
(10)
12

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
where Z (U or {(T")))> represents the latent features of the pockets molecules, which are the output of their respective encoders, and Tk is temperature scaling factor.
SimCLR contrastive learning: This contrastive learning task encourages the model to learn representations where similar pockets or molecules are mapped closer together in the feature space, while dissimilar pockets or molecules are pushed farther apart. This task enhances the model's ability to distinguish between different types of samples. Specifically, through the aforementioned reconstruction tasks, we obtain masked pockets { molecules {(=g) ) ) which include both masked types and corrupted coordinates. These original and masked versions are treated as two distinct views of a pocket or a molecule. The features for each view are generated by their respective encoders, denoted as 2 E Z = {( or
1-1 or {fxx ((m) ((r))²> } respectively. SimCLR aims minimize the distance between the original and masked versions of these two views f. and f (f!" and f(")) while maximizing the distance between them and all other views, f. and ff (f:" and fTr) where i # j. The SimCLR loss can be formulated follows:
{(fT")))= and 2n 2 = = (ft)
Laim = (s [ka.,)],
exp (sim (Len,)) / /Tam)
(11)
le = log Enjt exp (sim (ks,, / Tsim)
where sim(-) denotes the dot product and Tsim is temperature scaling factor.
B. The hyperparameter settings
The hyperparameter settings for the main branch are exactly the same as those used in DrugCLIP(Gao et al., 2024), and the hyperparameter settings for the auxiliary branch model are shown in Table 5. To ensure that the losses from different tasks are on the same scale, we multiply each loss by a weight that controls the scale of each respective loss.
Table 5. Hyperparameter settings for auxiliary branch model.
Hyperparameter
Value
Tkl
1.0
Taim
0.007
Layers of weights allocation MLP
2
Layers of cat multi-layers of encoder MLP Mask ratio
3
0.15
Corrupted distance
[-1,1]
Vocabulary size (atom types)
30 (molecule) / (pocket) 1.0
Gradient clip norm
Embedding dim of coarse- coarse-grained tasks
128
Embedding dim of fine-grained tasks
512
Control weight for Control weight for Ltype
1.0 (molecule) / 1.0 (pocket)
5 (molecule) / 1.0 (pocket)
Control weight for Ccoord Control weight for Ldist Control weight for Lxx Control weight for LNP Control weight for Lii Control weight for Leim
10 (molecule) 1.0 (pocket)
0.01 (molecule) / 0.01 (pocket)
0.01 (molecule) / /0001 (pocket)
0.01 (molecule) / 0.01 (pocket)
1000 (molecule) / 1000 (pocket)
0.01 (molecule) / 0.01 (pocket)
13

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
C. Pseudo code
Pseudo codes for the training and the testing process are listed Algorithm and Algorithm 2, respectively.
Algorithm Training stage
Require: {(f) ) )== pairs of molecules and pockets; a, learning rates
Initialize the network with pre-trained weights 0" O!
for batched pairs do
Conduct auxiliary tasks in molecule branch: C
Update parameters of auxiliary branch heads and encoder layer norm:
(ohorms @awar. + {ohorms Fauar - a VV
Conduct the auxiliary tasks in pocket branch: L
Update parameters of auxiliary branch heads and encoder layer norm:
a VV Update the primary task using the adapted parameters and update:
( hoorm ol J + hoom )uar]
Op Opri VV pprr
end for Return ppri, f ouz: 0%
Algorithm 2 Testing stage
Require: An interesting pocket and a set of candidate molecules M [m,, M2,.......... mn); mp learning rate Initialize the network with trained weights ppps
for batched molecules in M do
Conduct auxiliary tasks in molecule branch: Cauz
Update parameters of auxiliary branch heads and encoder layer norm:
{ Taux) f- {orrrm aux J V
end for
Conduct the auxiliary tasks in pocket branch: L
Update parameters of auxiliary branch heads and encoder layer norm:
(ohorm ] f- (ohorm Jbuz] V
Complete the primary task using the adapted parameters and compute scores:
Scores Similarity( (Fo (p), Fom (M))
Reload ppri to prevent accumulation of offsets
Return Predicted scores for each molecule
14

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
D. Additional ablation study
conduct ablation experiments on the selection of various auxiliary tasks, as shown in Table 6. Ltype and NN are closely related, while Leoords Laist and P are also tightly linked. The ablation results reveal that without Liype and Lxx. the model fails converge, as demonstrated by the last two rows. The other results indicate that the inclusion of each loss term is meaningful.
Each row of the ablation study is explained in detail follows:
w/o Lkt: Drug-TTA without the KI. loss
w/o Leim: Drug-TTr TTA without the SimCLR loss
w/o Liype Lx Drug-TTA without the masked atom type construction loss and atom feature normalization loss
wlo Looord fiist LNP: Drug-TIa, -TA without the corrupted coordinate reconstruction loss, distance reconstruction loss, and pairwise normalization loss
w/o Lesord Edist LNx Cnp: Drug-TTA without the corrupted coordinate reconstruction loss, distance reconstruction loss, atom feature normalization loss and pairwise feature normalization loss
w/o LNx LNP: Drug-TTA without the atom feature normalization loss and pairwise feature normalization loss
w/o Liype LNX LNP: Drug- TTA without the masked atom type construction loss, atom feature normalization loss and pairwise feature normalization loss
Table 6, Ablation study self-supervised tasks.
EF
Method
AUROC(%)
BEDROC(%)
0.50%
1%
5%
DrugCLIP(Gao et al., 2024)
80.93
50.52
38.07
31.89
10.66
Drug-TTA
93.16
82.82
57.50
54.04
16.88
w/o coarse- grained tasks
82.22
53.25
39.87
33.89
11.47
w/o fine-grained tasks w/o Lii
79.35
41.77
31.58
26.05
9.53
88.38
69.72
50.30
44.61
14.60
wlo Laim
83.71
54.76
41.59
35.09
11.60
w/o Ligpe LNX
86.85
68.28
49.77
44.46
14.03
wlo Lecord Edist LNP
83.12
64,22
46.38
41.05
13.69
w/o Lcoord Ediet LNX CNP
79.58
40.81
29.43
24.86
9.52
w/o LNX LNP w/o Leype LNX CNP
Furthermore, conduct ablation study the times of TTA fine-tuning for the encoder. We compare the performance of performing two adjustments to the molecule encoder before applying it to the primary task, as shown in Table 7. The results reveal that multiple adjustments yield inferior performance compared to single update.
Table 7. Ablation study TTA fine-tuning times.
EF
Method
AUROC(%)
BEDROC(%) (%)
0.50% %
1 1%
5%
DrigCLIP(Gao et al., 2024)
80.93
50.52
38.07
31.89
10.66
Drug-TTA
93.16
82.82
57.50
54.04
16.88
TTA 2times
84.04
71.76
54.97
47.13
14.03
15

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
E. Details of training and testing
During both training and testing, we utilize the DrugCLIP dataset, which is derived from the PDBBind dataset (Wang et al., 2005) and augmented through the identification of homologous proteins. To achieve zero-shot setup, we exclude any overlapping data from the benchmark dataset when evaluating on different datasets. In the training phase, we optimize the primary task using the A optimizer with learning rate of le-3 and batch size of 48, with acceleration provided by an NVIDIA A40 GPU. For optimizing the auxiliary branch, we use the SGD optimizer, setting the learning rate for the molecule branch at 1e-3 and the pocket branch at lc-4. During inference, we update only the auxiliary branch, continuing with the SGD optimizer; the learning rate for the molecule branch is set at 0.005, while the pocket branch's 's learning rate is 0.0001, and the batch size is increased to 64.
In the zero-shot inference process, DrugCLIP has not been evaluated on the AD, DEKOIS 2.0, and CASF-2016 benchmarks. The results on these three datasets are based our own inference, while results the other datasets, including those of the comparison methods, are sourced from the original DrugCLIP paper. We reassess the overlap of data between the DrugCIIP training set and these three benchmarks. It is found that there is no data overlap between AD and CASF-2016 6, so we use the weights provided by DrugCLIP for inference on these two benchmarks. For DEKOIS 2.0, we remove the overlapping data and retrained the DrugCLIP model using its original hyperparameters.
16

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
F. Evaluation metrics
BEDROC (Boltzmann-Enhann) Diserimination Receptor-Ligand Interactions): BEDROC is a metric frequently applied in virtual screening, especially in drug discovery. emphasizes the early-ranked molecules by incorporating an exponential weighting factor. The purpose of BEDROC is to evaluate the binding efficacy of molecules the target, focusing primarily the pocket-ligand interactions, while considering the relative importance of molecules ranked earlier in the screening process. BEDROCss is common variant of BEDROC where the top 2% of candidates contribute 80% of the BEDROC score. The formula for BEDROC is defined as:
Σ ii=1 exp
Za sinh ( cosh (Σ ala)
BEDROC,
+ exp (a(1 2n))"
(12)
Za
1-oxp( -a) exp(a/N)-)
cosh (g)
where a the parameter controlling the sensitivity of the metric the early ranks in the list (typically ranges from 100), N is the number of compounds being considered, r3 is the rank of the i-th compound in the list, and Za normalization constant.
EF (Enrichment Factor): The Enrichment Factor (EF) is another commonly used metric for assessing how well active compounds are enriched within the top-ranked molecules from virtual screening. A higher EF suggests that the top-ranked molccules are more likely to active. The EF calculated as follows:
NTB.
EF,
(13)
NTB,
where NTB, is the count of active molecules in the top a% and NTB, is the total number of active compounds.
ROC Enrichment Metric (RE): The ROC enrichment metric (RE) is used to measure the ratio of the true positive rate to the false positive rate (FPR) at a specific threshold. It is defined as:
11 TP
RE(*%) P FP.x
(14)
where n is the total number of compounds, TP refers the true positive compounds correctly identified active, P the total number of active compounds, and FP 2% is the number of false positives predicted at given threshold.
We add the comparison results between DTA and DrugCLIP on the RE metric across five tasks, as shown in Table 8. It can be observed that incorporating the TTA strategy significantly improves the performance.
Table 8 Additional results on RE metric.
RE
Dataset
Method
0.50%
%
2 2%
5%
DrugCLIP
73.97
41.79
23.68
11.16
DUD-E
Drug-TTA TTA
144.38
77.48
40.9
17.09
DrugCLIP
9.03
5.54
3.46
2.27
LIT-PCBA
Drug-TTA
77.38
44.04
23.78
10.74
DrugCLIP
52.44
33.07
20.63
10.57
AD
Drug-TTA
128.13
71.71
38.98
16.79
DrII
48.89
31.30
18.36
9.20
DEKOIS 2.0
DTAg-TTA
107.84
58.18
31.68
13.73
DrugCLIP
96.33
53.09
30.42
13.43
CASF-2016
Drug-TTA
141.85
73.78
39.07
16.41
17

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
G. Visualization
As shown in Figure 5, we visualize the features of active molecules (indicated in blue) and inactive molecules (indicated in grey) encoded by the top, middle, and last layers of the molecular encoder for two targets. We can observe that the features extracted by different layers represent different levels of molecular characteristics. Shallow-layer features, middle-layer features and deep-layer features exhibit significantly different distributions, suggesting that integrating features from all three layers is meaningful for generating loss weights.
Moreover, we visualize the feature distribution heatmap for additional molecules, as shown in Figure 6.
Top layer
Middle layer
Last layer
CP3A4
DHIT
Figure 5. Visualization of feature distributions for the targets CP3A4 and DHI on the top. middle and last layer of molecule encoder. The blue points denote positive molecule fcatures, and the gray points indicate negative molecule features.
Sample
Molecule weight
LogP
Rotatable bonds
Carhon atom count
Weights allocation
512.65
2.94
16
29
1.4000
1.1583
511.67
1.82
8
27
473.60
1.77
7
22
1.0845
1.
1.198
1.1173
1.0839
1.4111
1.4400
1.5099
2.3034
1.7846
1.3464
1.0833
1.4733
1.4458
1.4297
2.2355
1.2028
1.1297
18
496.74
4.52
13
1,4710
449.62
4.76
6
16
477.57
-0.99
6
28
1.0836
1.4633
1,4237
1.4779
2.3979
1.2223
1.1333
485.99
3,24
8
26
1.0000
1.4306
1.4439
1.4725
2.2427
1.2430
1.3744
10
401.47
2.28
4
17
1.1018
1.43/8
1.4643
1.49/4
2.04.0
1.2038
1.3481
407.49
1.17
5
18
1.4654
468.59
13
1.4549
A
2
2
2
2,
2
Figure 6. Weight allocation heatmap of 12 molecules in TTA and their molecular properties. The heatmap illustrates the weight allocation for molecules during the TTA process. The corresponding molecular properties include molecular weight, LogP, number rotatable bonds, and carbon atom count.
18

---

Drug-TTA: Test-Time Adaptation for Drug Virtual Screening
H. Inference time and computational burden
conduct additional experiments comparing the memory cost and inference time of Drug-TTA and DrugCLIP under the same conditions (i.e., on an RTX 3090 GPU with batch size of 64).
Compared DrI, Drug- -TAA increases memory usage from 1313 MiB to 1805 MiB for the molecule branch per batch and from 1232 MiB to 1627 MiB for the pocket branch (with single target). Despite this increase, our model can still maintain batch size of 64, ensuring training efficiency remains unaffected. Moreover, given that standard 24 GB (24,576 MiB) GPU is widely accessible to researchers, this additional memory consumption remains well within practical limits and does not bring computational bottleneck.
Memory Cost:
Inference Time: Based on the average inference time under our experimental setup on the DUD-E benchmark, Drug-TTA requires 2.1 days for virtual screening at practical scale (100 million molecules for a single target), compared to 0.8 day for DrugCLIP. While TTA does introduce additional computation time due to model adaptation during inference, this overhead insignificant in the context of the overall drug discovery timeline, which spans years (including both in silico and wet-lab experiments). More importantly, the substantial performance improvement brought by Drug-TTA significantly reduces the trial-and-error burden in wet-lab experiments, making the additional computational time worthwhile.
In practice, given the relatively low hardware cost our approach, inference time can be further reduced with minimal investment computational resources if needed. We will incorporate this discussion on computational cost and inference time in the final version of the paper. Once again, we sincerely appreciate your interest in the practical applicability of our method.
19


Reviews:
[{"review_id": "Hkt4WunPU0", "reviewer": "Reviewer_wTg7", "summary": "This paper introduces Drug-TTA, a novel test-time adaptation (TTA) approach for drug virtual screening that leverages multi-task meta-auxiliary learning to adapt the model to each test instance. Drug-TTA incorporates a large number of self-supervised auxiliary tasks into both training and inference processes and proposes an Auxiliary Loss Balance Module (ALBM) based on multi-scale features to dynamically adjust the weights of auxiliary tasks during inference. Extensive experiments demonstrate its effectiveness.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper and for your thoughtful feedback. We greatly appreciate your recognition of the novelty and effectiveness of our work, as well as your positive remarks on our writing and presentation. In response to your comments, we conduct additional experiments and provide further visualization analysis. Below are our point-by-point responses to your concerns.\n## W1&Q2 (part 1): Concern regarding additional computational overhead and inference time  \nWe appreciate your concern regarding computational overhead and inference time of Drug-TTA. Due to response length constraints, we kindly refer you to our response to Reviewer doCT(W1), where we provide a detailed analysis. \n## W2: More detailed analysis and explanation of Figure 5 \nThank you for carefully reviewing our appendix. As shown in Figure 2(b), our ALBM concatenates features from the molecule encoder's top, middle, and last layers and uses the concatenated features to calculate weights for each loss to balance the auxiliary losses. Figure 5 illustrates the feature space distribution of molecules at different encoder layers (top, middle, and last) during virtual screening across different targets. Our visualization highlights the diversity in molecular feature representations at different layers, demonstrating the necessity of fusing features from the three layers in ALBM. This observation is also consistent with the results of our ablation study in Table 4, which demonstrates that fusing features from all three layers is essential for optimal performance.\n## W3: More detailed analysis of ablation study results  \nThank you for acknowledging the substantial workload of our ablation studies. We understand your concern regarding the significant performance drop when removing certain components. This is because the key components of Drug-TTA are highly interdependent, with each component playing a crucial role. Below, we provide a detailed explanation:  \n- **Auxiliary Branch Selection**: Performing TTA on only one branch (either molecule or pocket) leads to a mismatch in feature adaptation, as the unadapted branch struggles to align within the shared feature space. Thus, TTA is necessary for both the molecule and pocket branches.  \n- **Multi-Scale Feature Layer Selection**: Our framework is designed to comprehensively capture molecular features across multiple layers. While removing any single layer leads to a performance drop, the model still outperforms DrugCLIP, demonstrating the robustness of our multi-scale feature design.  \n- **Self-Supervised Task Selection**: Self-supervised tasks are crucial for learning both fine-grained and coarse-grained molecular representations, where coarse-grained features aggregate fine-grained representations. Removing the fine-grained feature learning task results in suboptimal performance improvement.  \n- **Weight Regularization & Meta-Learning**: Weight regularization prevents model collapse during training. Meta-learning is the core mechanism that prevents overfitting to auxiliary tasks. Removing these components leads to a significant performance drop, highlighting their necessity.\n\nOverall, most ablation settings still outperform DrugCLIP, validating the effectiveness of our design. We assure you that upon acceptance, we will release our original code and model weights to facilitate reproducibility.  \n## W4: Performance analysis on LIT-PCBA  \nWe appreciate your concern regarding the performance analysis on LIT-PCBA. Due to response length constraints, we kindly refer you to our response to Reviewer neHb(W3), where we provide a detailed analysis.    \n## Q1: How did you consider adjusting the normalization layers in the encoder using auxiliary tasks during test time? What would happen if all parameters were fine-tuned?\nThank you for your question. We conduct an additional experiment with full fine-tuning, and the results on the DUD-E dataset are shown below. Fine-tuning all parameters with limited data leads to instability due to the large parameter space and distribution shift, causing gradient explosion or model collapse. This ultimately degrades performance, as reflected in our results. Therefore, we consider adjusting the normalization layers, as in other TTA methods.\n|Method|AUROC(%)|BEDROC(%)|EF 0.50%|EF 1%|EF 5%| \n|-|-|-|-|-|-|\n|**DrugCLIP**|80.93|50.52|38.07|31.89|10.66| \n|**Drug-TTA**|93.16|82.82|57.50|54.04|16.88| \n|**All parameters fine-tuned**|64.34|11.32|7.20|5.97|3.81|\n## Q2(Part 2): Is it necessary for the loss of auxiliary task to converge?\nThank you for your question. Our method does not require the auxiliary task loss to converge, as the model only needs to perform a single adjustment during testing. This aligns with the common practice in TTA, where most methods apply only one or a few updates rather than waiting for full convergence. This design ensures efficiency and generalizability, preventing excessive computation and overfitting to individual test samples.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "This paper does not include much theoretical part.", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "Refer to the strengths and weakness.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths:\n\n1. The proposed method significantly outperforms the baseline DrugCLIP on multiple datasets.\n2. The experimental section is comprehensive.\n3. The writing is clear, particularly with the well-illustrated framework diagram.\n\n\nWeakness:\n\n1. Parameter fine-tuning is required during the TTA (Test-Time Adaptation) phase, and new targets all need to be re-fine-tuned, which may incur additional computational overhead.\n2. I'm a bit puzzled by the analysis of Figure 5. Could the authors provide a more detailed analysis and explanation?\n3. It seems that there are numerous hyperparameters to design and many tricks involved. Although the authors conducted ablation studies on almost every module, I found that the performance drops significantly when one setting is missing, which is somewhat puzzling. On one hand, this indicates that the workload is substantial. However, on the other hand, it may increase the difficulty of reproducing the results.\n4. As is well known, the Lit-PCBA dataset is relatively challenging. Why does the proposed method achieve such significant performance improvements on the Lit-PCBA dataset compared to other baselines, and even surpass the performance on simpler tasks like DUD-E? I believe more analysis is needed to explain why the method works so well in this context.", "other_comments_or_suggestions": "NA", "questions_for_authors": "I have a few detailed questions:\n\n\n1. How did you consider adjusting the normalization layers in the encoder using auxiliary tasks during test time? What would happen if all parameters were fine-tuned?\n2. How long does it take to fine-tune during the test phase, and is it necessary for the loss of the auxiliary task to converge?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g9wEGaXwpx", "reviewer": "Reviewer_neHb", "summary": "The paper introduces Drug-TTA, a novel test-time adaptation (TTA) framework for drug virtual screening that dynamically adjusts a pre-trained model to each test instance. Unlike traditional zero-shot screening methods, Drug-TTA utilizes self-supervised auxiliary tasks to adapt its protein and molecule encoders at inference time, enhancing generalization to unseen data. A key innovation is the multi-scale feature-based Auxiliary Loss Balance Module (ALBM), which optimally balances auxiliary task losses. Additionally, meta-learning ensures that adaptation enhances rather than hinders the primary task. Drug-TTA significantly outperforms state-of-the-art (SOTA) methods across five benchmarks (DUD-E, LIT-PCBA, AD, DEKOIS 2.0, CASF-2016), with an average AUROC improvement of 9.86%, demonstrating its effectiveness in zero-shot-like virtual screening while not being strictly zero-shot due to test-time updates.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper and for your thoughtful feedback. We greatly appreciate your recognition of our innovation and the practical significance of our work. Below, we provide detailed responses to your concerns.  \n## W1: Concern regarding \"Zero-Shot\" strict definition  \nWe sincerely appreciate your comment. You are correct that our method involves updating model parameters during testing, which does not align with the definition of zero-shot learning. Nevertheless, as you commented, our method does work in a zero-shot-like setting. Because of the very similar working condition, some recent TTA [1-3] works also define methods like ours, where no labeled samples are used during testing, as zero-shot. We will clarify the difference in revision.  \n\nWe acknowledge that our comparison methods do not adjust model parameters at test time, which may raise concerns about fairness. However, our input conditions remain identical across all methods, as we do not use any labeled test samples. Moreover, since our work is the first to introduce TTA in this task, there are no directly comparable baselines. Given these constraints, we compare against the most relevant existing approaches. We appreciate your insightful remark and will refine the description of zero-shot learning in the final version of our paper.  \n### References:  \n[1] Liberatori B, Conti A, Rota P, et al. Test-time zero-shot temporal action localization. CVPR 2024.  \n[2] Zhao S, Wang X, Zhu L, et al. Test-Time Adaptation with CLIP Reward for Zero-Shot Generalization in Vision-Language Models. ICLR 2024.  \n[3] Aleem S, Wang F, Maniparambil M, et al. Test-time adaptation with SALIP: A cascade of SAM and CLIP for zero-shot medical image segmentation. CVPR 2024.  \n## W2: Concern regarding inference time and computational burden  \nWe appreciate your concern regarding inference time and computational burden of our method. Due to response length constraints, we kindly refer you to our response to Reviewer doCT(W1), where we provide a detailed analysis.\n## W3: Performance analysis on LIT-PCBA  \nThank you for carefully reviewing our experimental results and for your valuable suggestions. We supplement our analysis with additional experimental and visualization results as follows.  \n\nLIT-PCBA is a highly challenging task due to its extreme class imbalance, with only 0.74% of the screened molecules being active. This makes all methods have much lower performance on this dataset than on the other benchmarks and leaves much room for improvement. Drug-TTA, with its ability to adapt to each sample, is particularly well-suited to handling such challenging scenarios.\n  \nSpecifically, in the figure of https://pasteboard.co/Go21VXS5RV5H.bmp, we visualize the feature distributions for the MTORC1 target on the LIT-PCBA benchmark using both Drug-TTA and DrugCLIP. The visualization clearly shows that positive molecules are positioned closer to the pocket features in Drug-TTA, resulting in higher ranking scores for these molecules. Consequently, this enhances the early retrieval capability, which explains the substantial improvements in BEDROC and EF—metrics that are highly sensitive to early ranking performance.  \nFurthermore, we summarize the performance improvements observed for several targets:  \n| Target  | Active Molecule Proportion | AUROC (Drug-TTA) | BEDROC (Drug-TTA) | AUROC (DrugCLIP) | BEDROC (DrugCLIP) |\n|-|-|-|-|-|-|\n| **MTORC1** | 0.29% | 98.81% | 48.38% | 67.38% | 1.74% |\n| **PKM2**   | 0.22% | 93.71% | 59.31% | 74.28% | 0.90% |\n| **FEN1**   | 0.10% | 98.35% | 58.36% | 87.51% | 1.79% |\n\nFrom the table, it is evident that Drug-TTA is more robust to extreme class imbalance. This adaptation significantly enhances early retrieval performance.  \n\nWe must **clarify that** our performance on LIT-PCBA only surpasses the performance on DUD-E in one metric, EF (0.50%). This is because EF focuses solely on the top-ranked molecules, and given the low proportion of active molecules in the dataset, EF can be more easily improved as active molecules become more concentrated at the top. However, regarding AUROC and BEDROC, which are influenced by the overall ranking, the performance on LIT-PCBA is lower compared to other benchmarks.  \n\n## Other Comments or Suggestions: Table captions need improvement  \nThank you for your comment on the table caption. Since we are unable to modify the original manuscript at this stage, we will make the revisions in the final version of the paper. Specifically, we plan to add the following clarifications: In Table 123: \"AUROC, BEDROC, and EF are reported (higher values indicate better performance). Bold values represent the best performance, and green indicates improvements of Drug-TTA over DrugCLIP.\" In Table 4: \"*\"w/o\"* indicates the removal of the respective component compared to Drug-TTA.\"", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Not available.", "experimental_designs_or_analyses": "Yes, but there seems to be a difference in evaluation settings with comparison baselines. The baseline models are zero-shot settings, but strictly speaking, Drug-TTA does not seem to be a zero-shot setting (due to Algorithm 2).", "relation_to_broader_scientific_literature": "In the case of molecules, it is difficult to find good negative samples, and we want to address this through test-time adaptation. This seems to have high utility not only in virtual screening but also in other molecular property prediction fields.", "essential_references_not_discussed": "Related works are properly discussed.", "other_strengths_and_weaknesses": "[Strengths]\n\n- Incorporating five self-supervised auxiliary tasks (both fine-grained and coarse-grained) allows the model to learn useful representations without labeled binding affinity data, making it adaptable and efficient.\n\n- Drug-TTA outperforms existing methods across five major virtual screening benchmarks, achieving an average AUROC improvement of 9.86% over DrugCLIP, demonstrating its superior accuracy in ranking active molecules.\n\n[Weaknesses]\n\n- While it operates in a zero-shot setting, Drug-TTA modifies model parameters during test-time inference, which contradicts the strict definition of zero-shot learning, making it less interpretable in direct zero-shot comparisons.\n\n- Test-time adaptation requires multiple forward and backward passes per test instance, increasing inference time and computational burden compared to traditional zero-shot methods like DrugCLIP.\n\n- The paper reports a significant improvement in performance on the LIT-PCBA benchmark, particularly in EF (Enrichment Factor) metrics, where Drug-TTA achieves an order-of-magnitude improvement over DrugCLIP. However, the paper does not provide a thorough discussion on why Drug-TTA performs exceptionally well on LIT-PCBA compared to other benchmarks.", "other_comments_or_suggestions": "The table captions are poor and need to be improved.", "questions_for_authors": "Please see the weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kKMAhFudQf", "reviewer": "Reviewer_doCT", "summary": "The authors pinpoint two issues in ML-based structure-based drug discovery: a lack of negative pairs for ML-based docking methods and overly simplistic negative pairs in contrastive learning approaches (e.g., DrugClip), leading to a domain shift during inference when most screened molecules are inactive. \nTo mitigate this, they propose a test-time adaptation pipeline centered on DrugClip that updates the encoders using multiple auxiliary tasks, where the weight of each task is adjusted per input, and using meta-learning to prevent overfitting to these auxiliary tasks. \nThe authors show that their methods, Drug-TTA, significantly improved over their baseline method DrugCLIP and is the new SOTA.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper and for your thoughtful and thorough feedback. We truly appreciate your recognition that the evidence for the claims is clear and convincing, as well as your positive comments on our writing and figures. Below, we will address your questions and concerns in detail.\n## Methods and Evaluation Criteria: additional comparison method\nThank you for the suggestion. We will discuss this method and cite the new arXiv paper in the final version. While our method does not surpass SPRINT [1] in AUROC on LIT-PCBA, it achieves improvements in BEDROC and EF. We will incorporate this comparison accordingly.\nPerformance Comparison:\n|Method| AUROC(%)| BEDROC(%)|EF 0.50% |EF 1%|EF 5%|\n|-|-|-|-|-|-|\n|DrugCLIP|57.17|6.23| 8.56 | 5.51| 2.27|\n|Drug-TTA|71.24| **45.08**| **74.39**|**42.74**|**10.61**|\n|SPRINT[1]| **73.4**|12.3|15.90|10.78|5.92|\n\n[1] [McNutt'24] SPRINT Enables Interpretable and Ultra-Fast Virtual Screening against Thousands of Proteomes, http://arxiv.org/abs/2411.15418\n## Experimental Designs or Analyses: replicates on all datasets\nThank you very much for raising this question. We conduct additional testing of our method on five benchmarks, and the results are averaged over three repetitions. Below are the mean and standard deviation for each benchmark:\n|Benchmark| AUROC(%)| BEDROC(%)| EF 0.50%| EF 1%| EF 5%|\n|-|-|-|-|-|-|\n|DUD-E|93.15±0.01|82.72±0.27|57.47±0.09|53.90±0.13|16.85±0.06|\n|LIT-PCBA|71.19±0.19|44.87±0.06|74.26±0.24|42.44±0.37|10.43±0.03|\n|AD|92.68±0.08|86.80±0.10|32.56±0.03| 30.56±0.06|15.31±0.03|\n|DEKOIS2.0|83.43±0.24|73.30±0.31|26.63±0.28|25.43±0.21|12.94±0.11|\n|CASF-2016|91.82±0.08|86.01±0.25|42.84±0.12|41.49±0.19|16.05±0.02|\n\n## W1: Concern regarding inference time and computational burden\nThank you for raising this question. We apologize for not providing a detailed explanation in the paper. To address this issue, we conduct additional experiments comparing the memory cost and inference time of Drug-TTA and DrugCLIP under the same conditions (i.e., on an RTX 3090 GPU with a batch size of 64).  \n- **Memory Cost**: Compared to DrugCLIP, Drug-TTA increases memory usage from 1313 MiB to 1805 MiB for the molecule branch per batch and from 1232 MiB to 1627 MiB for the pocket branch (with a single target). Despite this increase, our model can still maintain a batch size of 64, ensuring training efficiency remains unaffected. Moreover, given that a standard 24 GB (24,576 MiB) GPU is widely accessible to researchers, this additional memory consumption remains well within practical limits and does not bring computational bottleneck. \n- **Inference Time**: Based on the average inference time under our experimental setup on the DUD-E benchmark, Drug-TTA requires 2.1 days for virtual screening at a practical scale (100 million molecules for a single target), compared to 0.8 day for DrugCLIP. While TTA does introduce additional computation time due to model adaptation during inference, this overhead is insignificant in the context of the overall drug discovery timeline, which spans years (including both in silico and wet-lab experiments). More importantly, the substantial performance improvement brought by Drug-TTA significantly reduces the trial-and-error burden in wet-lab experiments, making the additional computational time worthwhile.  \n\nIn practice, given the relatively low hardware cost of our approach, inference time can be further reduced with minimal investment in computational resources if needed. We will incorporate this discussion on computational cost and inference time in the final version of the paper. Once again, we sincerely appreciate your interest in the practical applicability of our method.  \n## C1: Redrawing Figure 3\nThank you for your suggestion regarding our visualizations. We redraw Figure 3 with an increased size of the pocket representation, improved color schemes, and clearer markers for better readability. Since images cannot be uploaded here, please refer to the updated figure at https://pasteboard.co/Up1mGfvCKsYN.bmp.\n## C2: Clarification on structure-based virtual screening\nThank you for your suggestion. We will clarify that we are talking about *structure-based* virtual screening in the Introduction.\n## C3: Value size and decimal explanation\nWe appreciate your feedback. We redraw Figure 4 and enlarge the values in the middle heatmap, as shown in https://pasteboard.co/chAPEluE4BNN.bmp. The four-decimal precision is chosen to better show the differences in sample weights.\n## Q1: Clarification of framework diagram\nThank you for your thorough review. We sincerely apologize for any misunderstanding caused by the design of our main framework diagram. We will clarify this in the paper. The common MLP is just employed to project the concatenated three-layer features (512 × 3 = 1536) into 128 dimensions. The subsequent seven MLPs are specifically designed to accommodate seven auxiliary losses. We will update our description in the Method.", "claims_and_evidence": "Evidence for the claims is clear and convincing.", "methods_and_evaluation_criteria": "The proposed method was evaluated on established virtual screening benchmarks using standard metrics, and compared against both well-established and previous state-of-the-art approaches. \nHowever, the evaluation does not include a comparison to the recent SPRINT [McNutt'24] method, which is, to my knowledge, the current SOTA for LIT-PCBA.\n\n[McNutt'24] SPRINT Enables Interpretable and Ultra-Fast Virtual Screening against Thousands of Proteomes, http://arxiv.org/abs/2411.15418", "theoretical_claims": "Application paper, doesn't contain proof. \nIn general, didn't check the correctness of the formulas.", "experimental_designs_or_analyses": "- In general, the experimental design is sound using and their analysis seems sound.\n- Furthermore, the proposed methods is quite complex with many components and the authors have done a systematic ablation study on each of the components.\n- However, not a single value in the paper comes with an error bar giving limited significance to the results. \n   I'm very aware this is systematic problem with the community, and these error bars are never reported on these VS benchmarks, but you can do better. \n  I'm not asking you to produce triplicates for all of your experiments, but I would like to see replicates and the analysis for at least one of the datasets (e.g., CASF-2016).", "relation_to_broader_scientific_literature": "Building on prior work in virtual screening (VS), test-time adaptation (TTA), and meta-learning, the proposed approach, Drug-TTA, enhances generalizability by: (1) introducing TTA to drug virtual screening, and (2) achieving state-of-the-art performance on several VS benchmarks.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "- S1: Paper is generally well-written.\n- S3: Figure 1 very nice and gives an excellent overview of the method! \n- W1: With virtual screening libraries like GDB-17 and Enamine Real now containing billions of compounds, the efficiency of VS methods is more critical than ever. \n   While Drug-TTA shows promise, it introduces additional complexity over DrugClip by incorporating extra parameters and performing parameter updates during inference. \n   A direct comparison of parameter counts and inference times between DrugClip and Drug-TTA would be invaluable for assessing the practical efficiency of the proposed method.", "other_comments_or_suggestions": "- C1: In figure 3, finding the pocket representation and the contrast between the actives and inactives is not very clear. \n   Maybe increase the size of pocket representation dot and change the color scheme.\n- C2: In the introduction maybe clarify that your are talking about *structure-based* virtual screening.\n- C3: Values in the Fig 4. heatmap are way too small. \n   Also is 4 decimals really required in a illustration?", "questions_for_authors": "- Q1: As I understand the text and equation in 3.2, the concatenated features directly are passed to the seven different MLPs, but in Figure 1 it seems that the concatenated features are first passed in to a common MLP before having a separate MLP for each task. \n   Can you clarify on this?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lJsfUSuj4x", "reviewer": "Reviewer_53jo", "summary": "The paper introduces Drug-TTA, a novel framework for drug virtual screening that incorporates test-time adaptation through multi-task meta-auxiliary learning. The authors build upon a contrastive learning paradigm by integrating a series of self-supervised auxiliary tasks-three fine-grained tasks (masked atom type prediction, corrupted coordinate reconstruction, and distance prediction), and two coarse-grained tasks (KL divergence minimization and SimCLR contrastive learning)-to dynamically adjust the model during inference. Central to their approach is the multi-scale feature-based Auxiliary Loss Balance Module (ALBM), which adaptively computes loss weights for each auxiliary task based on features extracted from different layers of the encoder. Moreover, a meta-learning strategy is employed to harmonize the auxiliary tasks with the primary virtual screening objective, ensuring that the test-time adaptations improve the main task without causing overfitting. Extensive experiments across five zero-shot virtual screening benchmarks demonstrate that Drug-TTA significantly outperforms existing methods, particularly showing notable improvements in AUROC, BEDROC, and enrichment factor metrics.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper and for your thoughtful comments. We sincerely appreciate your recognition of the novelty and significant performance of our method. Below, we provide detailed responses to your questions and concerns.\n\n## W1: Losses $L_{NX}$ and $L_{NP}$ explanation\nWe sincerely apologize for the insufficient explanation of $L_{NX}$ and $L_{NP}$. Below, we provide a more detailed clarification of these loss terms.\nOur encoder follows the overall framework of Uni-Mol, and the design of  $L_{NX}$ and $L_{NP}$ is aligned with its pretraining tasks. Specifically, $L_{NX}$ is introduced to normalize masked atom features. Given a molecular representation $X \\in \\mathbb{R}^{B \\times T \\times d}$, where $B$ is the batch size, $T$ is the number of atoms, and $d$ is the feature dimension, the norm constraint is formulated as:\n\n$\n\\mathcal{L}_{NX} = \\max(0, ||x||_2 - \\sqrt{d} - \\tau),\n$\nwhere $||x||_2$ ensures numerical stability, and $\\tau$ is a tolerance margin. This constraint prevents masked atom features from deviating excessively from the expected norm, thereby stabilizing feature learning and preserving consistency across different masked regions.\n\nSimilarly, $L_{NP}$ is applied to pairwise relational representations, ensuring that the learned attention patterns do not fluctuate excessively when certain atoms are masked. This is crucial because molecular graphs inherently depend on structural constraints, and unnormalized pairwise features may lead to unstable or biased attention distributions.\n\nBy enforcing these norm constraints, our approach ensures that both atomic-level and relational representations remain well-calibrated, leading to more robust and generalizable molecular embeddings. We will revise the manuscript to incorporate this explanation and appreciate your insightful comments.\n\n## W2: Table 6 revision\n\nThank you very much for carefully reviewing our Appendix. We sincerely apologize for any lack of clarity in our original statement. Our revised version is as follows:\n\n- Drug-TTA without coarse-grained tasks\n- Drug-TTA without fine-grained tasks\n- Drug-TTA without the KL loss\n- Drug-TTA without the SimCLR loss\n- Drug-TTA without the masked atom type construction loss and atom feature normalization loss\n- Drug-TTA without the corrupted coordinate reconstruction loss, distance reconstruction loss, and pairwise normalization loss\n- Drug-TTA without the corrupted coordinate reconstruction loss, distance reconstruction loss, atom feature normalization loss and pairwise feature normalization loss\n- Drug-TTA without the atom feature normalization loss and pairwise feature normalization loss\n- Drug-TTA without the masked atom type construction loss, atom feature normalization loss and pairwise feature normalization loss\n\n## W3: Concern regarding computational efficiency, scalability, or model complexity\nThank you for raising this question. We apologize for not providing a detailed explanation in the paper. To address this issue, we conduct additional experiments comparing the memory cost and inference time of Drug-TTA and DrugCLIP under the same conditions (i.e., on an RTX 3090 GPU with a batch size of 64).  \n- **Memory Cost**: Compared to DrugCLIP, Drug-TTA increases memory usage from 1313 MiB to 1805 MiB for the molecule branch per batch and from 1232 MiB to 1627 MiB for the pocket branch (with a single target). Despite this increase, our model can still maintain a batch size of 64, ensuring training efficiency remains unaffected. Moreover, given that a standard 24 GB (24,576 MiB) GPU is widely accessible to researchers, this additional memory consumption remains well within practical limits and does not bring computational bottleneck. \n- **Inference Time**: Based on the average inference time under our experimental setup on the DUD-E benchmark, Drug-TTA requires 2.1 days for virtual screening at a practical scale (100 million molecules for a single target), compared to 0.8 days for DrugCLIP. While TTA does introduce additional computation time due to model adaptation during inference, this overhead is insignificant in the context of the overall drug discovery timeline, which spans years (including both in silico and wet-lab experiments).  More importantly, the substantial performance improvement brought by Drug-TTA significantly reduces the trial-and-error burden in wet-lab experiments, making the additional computational time worthwhile.  \n\nIn practice, given the relatively low hardware cost of our approach, inference time can be further reduced with minimal investment in computational resources if needed.  We will incorporate this discussion on computational cost and inference time in the final version of the paper. Once again, we sincerely appreciate your interest in the practical applicability of our method.  \n## W4: Code release\nWe assure you that upon acceptance, we will release our original code and model weights to facilitate reproducibility.", "claims_and_evidence": "The claims are supported by experimental results showing improvements in AUROC, BEDROC, and enrichment factors across multiple benchmarks. However, the paper lacks efficiency analysis, which is crucial for assessing practicality, particularly in large-scale applications. Additionally, the generalization of the approach across diverse datasets and real-world scenarios is not sufficiently validated, weakening the overall claims.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria in the paper are appropriate for the problem at hand.\nThe evaluation criteria, including performance metrics such as AUROC, BEDROC, and enrichment factor, are standard and relevant for assessing the effectiveness of virtual screening methods.", "theoretical_claims": "The paper does not include any formal proofs or rigorous theoretical claims that would require verification.", "experimental_designs_or_analyses": "The experimental designs and analyses appear sound, with a clear comparison of Drug-TTA against baseline methods on multiple benchmark datasets. The metrics used, including AUROC, BEDROC, and enrichment factor, are standard and relevant for evaluating virtual screening methods. \nHowever, there is no detailed analysis of the method's efficiency or scalability, which leaves questions about its practical deployment in large-scale settings unanswered.", "relation_to_broader_scientific_literature": "The Drug-TTA approach proposed in this paper echoes existing scientific literature in several aspects. First, the DrugCLIP framework significantly improves the performance of virtual screening by redefining virtual screening as a dense search task and using contrastive learning to align protein pockets and molecular representations.（https://arxiv.org/abs/2310.06367?utm_source=chatgpt.com）\n\nIn addition, the Point-TTA method adopts a test-time adaptation strategy in the point cloud registration task and designs three self-supervised auxiliary tasks to adapt to the new data distribution in the testing phase.（https://arxiv.org/abs/2308.16481?utm_source=chatgpt.com）\n\nThese studies provide theoretical foundations and practical experience for the Drug-TTA method, demonstrating the effectiveness of adaptive and contrastive learning strategies when applying tests in different tasks.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths:\nS1. The paper presents an innovative integration of test-time adaptation with multi-task meta-auxiliary learning, creatively combining contrastive learning with several self-supervised tasks. The proposed method addresses the domain shift between training and test data, which is a significant issue in virtual screening. \nS2. It demonstrates significant performance improvements on multiple virtual screening benchmarks.\nS3.The framework builds on recent advances (e.g., DrugCLIP, TTA) and addresses the challenging zero-shot setting in drug screening.\n\nWeaknesses：\nW1. The losses L_NX  and L_NP  appear in the framework diagram and loss function but are not explicitly introduced in the methodology. This omission makes it difficult for readers without prior knowledge to fully understand their role in the model.\nW2. The ablation study table, table 6, does not clearly define what each experiment represents. For example, “w/o-L_kl” should be explicitly noted as “Drug-TTA without the KL loss” to improve clarity.\nW3. The paper does not include any analysis of computational efficiency, scalability, or model complexity, which are crucial for assessing real-world applicability.\nW4. The authors do not provide code or implementation details, making it difficult to verify whether the method can be reproduced in practice.", "other_comments_or_suggestions": "None", "questions_for_authors": "Please address the concerns raised in W1–W4.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces Drug-TTA, a novel test-time adaptation (TTA) method for drug virtual screening. The approach leverages self-supervised auxiliary tasks and meta-learning to dynamically adapt the model to each test instance, aiming to improve the performance of virtual screening tasks.


Based on the reviews and rebuttal discussions:
- The paper presents a novel, well-executed methodology with significant performance gains.
- Key concerns about computational burden and hyperparameter fine-tuning were acknowledged but sufficiently addressed.
- While there remain concerns about efficiency, the authors demonstrated clear advantages over existing methods.

Given these considerations, it is a boarderline paper. The recommendation leans toward acceptance, if there is room in the program.
Decision: Accept (poster)