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

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.

W1: Losses $L_{NX}$ and $L_{NP}$ explanation

We sincerely apologize for the insufficient explanation of $L_{NX}$ and $L_{NP}$. Below, we provide a more detailed clarification of these loss terms. Our 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:

$\mathcal{L}_{NX} = \max(0, ||x||_2 - \sqrt{d} - \tau),$ where $||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.

Similarly, $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.

By 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.

W2: Table 6 revision

Thank 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:

• Drug-TTA without coarse-grained tasks
• Drug-TTA without fine-grained tasks
• Drug-TTA without the KL loss
• Drug-TTA without the SimCLR loss
• Drug-TTA without the masked atom type construction loss and atom feature normalization loss
• Drug-TTA without the corrupted coordinate reconstruction loss, distance reconstruction loss, and pairwise normalization loss
• Drug-TTA without the corrupted coordinate reconstruction loss, distance reconstruction loss, atom feature normalization loss and pairwise feature normalization loss
• Drug-TTA without the atom feature normalization loss and pairwise feature normalization loss
• Drug-TTA without the masked atom type construction loss, atom feature normalization loss and pairwise feature normalization loss

W3: Concern regarding computational efficiency, scalability, or model complexity

Thank 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).

• 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.
• 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.

In 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.

W4: Code release

We assure you that upon acceptance, we will release our original code and model weights to facilitate reproducibility.

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.

Methods and Evaluation Criteria: additional comparison method

Thank 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. Performance Comparison:

[1] [McNutt'24] SPRINT Enables Interpretable and Ultra-Fast Virtual Screening against Thousands of Proteomes, http://arxiv.org/abs/2411.15418

Experimental Designs or Analyses: replicates on all datasets

Thank 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:

W1: Concern regarding inference time and computational burden

Thank 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).

• 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.
• 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.

In 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.

C1: Redrawing Figure 3

Thank 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.

C2: Clarification on structure-based virtual screening

Thank you for your suggestion. We will clarify that we are talking about structure-based virtual screening in the Introduction.

C3: Value size and decimal explanation

We 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.

Q1: Clarification of framework diagram

Thank 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.

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.

W1: Concern regarding "Zero-Shot" strict definition

We 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.

We 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.

References:

[1] Liberatori B, Conti A, Rota P, et al. Test-time zero-shot temporal action localization. CVPR 2024.
[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.
[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.

W2: Concern regarding inference time and computational burden

We 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.

W3: Performance analysis on LIT-PCBA

Thank you for carefully reviewing our experimental results and for your valuable suggestions. We supplement our analysis with additional experimental and visualization results as follows.

LIT-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.

Specifically, 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.
Furthermore, we summarize the performance improvements observed for several targets:

From the table, it is evident that Drug-TTA is more robust to extreme class imbalance. This adaptation significantly enhances early retrieval performance.

We 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.

Other Comments or Suggestions: Table captions need improvement

Thank 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."

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.

W1&Q2 (part 1): Concern regarding additional computational overhead and inference time

We 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.

W2: More detailed analysis and explanation of Figure 5

Thank 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.

W3: More detailed analysis of ablation study results

Thank 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:

• 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.
• 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.
• 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.
• 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.

Overall, 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.

W4: Performance analysis on LIT-PCBA

We 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.

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?

Thank 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.

Q2(Part 2): Is it necessary for the loss of auxiliary task to converge?

Thank 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.