UniMatch: Universal Matching from Atom to Task for Few-Shot Drug Discovery

This work presents UniMatch for molecular property prediction in the few-shot learning scenario. It combines explicit hierarchical molecular matching with implicit task-level matching via meta-learning. The effectiveness of the few-shot learning and multi-task generalization is validated on the MoleculeNet, FS-Mol, and Meta-MolNet benchmarks. An ablation study and visualizations are provided to demonstrate the importance of hierarchical molecular matching and task-level matching.

This study introduces a new framework for molecular representation specially aimed at multi-task few-shot learning. In this context for each task there is a small support subset of molecules with known labels ($S_T$) and a query subset ($Q_T$) with unknown labels.

The contributions of the study can be broken down into the following components:

1. Matching at different levels of the molecular graph. This is achieved by an attention mechanism where the keys are the molecular features of the $S_T$, the molecular representations of $Q_T$ are used as query, and as value, the labels of the $S_T$. This matching attention mechanism is not novel. The novelty resides on repeating this matching with the representations at each layer of the GNN (thus effectively capturing information at different structural levels of any given molecule) and concatenating the resulting values that are then passed through an MLP to predict the final target value.
2. Meta-learning. An additional implicit matching mechanism is used to learn the similarity between different tasks and how to leverage the multi-scale representations from one task to another. This is achieved by an inner loop where the task-specific adaption parameters are updated using the matrix that captures the similarity between tasks and an outer loop where this similarity matrix is updated.

The results obtained in benchmark datasets clearly show that the benefits of this approach when compared to current state-of-the-art alternatives. The ablation study also clearly supports the contribution of all the components of the method.

The paper presents UniMatch, a framework designed to address data scarcity in few-shot drug discovery. It uses a dual matching approach: explicit hierarchical molecular matching, which captures information from the atomic level to full molecular structures, and implicit task-level matching via meta-learning to enable effective generalization across tasks.

The authors propose UniMatch, a dual matching framework that integrates hierarchical molecular matching with meta-learning, significantly improves the performance and generalization in molecular property prediction tasks.

Dear Reviewers and ACs,

As the discussion is ending in the next few hours, we are grateful for the insightful comments and valuable suggestions from all reviewers. In the following, we would like to summarize the contributions and responses of this paper again.

Contributions:

• To tackle molecular property prediction in few-shot scenarios, we propose a unified matching framework that combines hierarchical molecular matching with task-level matching through meta-learning. This approach is effective (R XF6i), innovative (R XF6i), promising (R PgJj), and adaptable to various domains with high social value (R 6vhN).
• The experiments are extensive (R PgJj, R 6vhN, R E3W3), covering three benchmarks, and the ablation study is complete (R 6vhN).
• The paper is well-written (R XF6i, R PgJj, R E3W3), and the figures are well-crafted (R XF6i, R E3W3).

Responses and Revisions:

• For Reviewer XF6i's concerns:

  • Q1 Further improving model interpretability: We analyzed the matching scores and parameters of each layer, showing their complementary roles in decision-making. We also plan to further enhance interpretability using DeepLIFT and the Kolmogorov-Arnold Network (KAN).

  • Q2 AUCPR and AUPRC: To avoid confusion, we have standardized the terminology to "AUPRC" throughout the paper.

• For Reviewer PgJj's concerns:

  • Q1 Mean pooling may result in the loss of details: We compared mean pooling with attention-based and ensemble-based pooling. The latter methods perform better at preserving fine-grained details. We will incorporate these advanced pooling strategies in future work.

  • Q2 Length of the related work section: We agree with the suggestion to condense the related work section and move it to the end of the manuscript. This change will improve the paper’s flow and readability.

• For Reviewer 6vhN's concerns:

  • Q1 More QM tasks in MoleculeNet: We’ve added QM9 experiments to our evaluation, where UniMatch achieves the best performance with the lowest MAE.

  • Q2 More idea about Inductive bias in attention: We plan to incorporate structural priors at different levels (atomic, substructural, molecular) to guide attention and improve performance.

  • Q3 Clarifying dual matching mechanism: The dual matching mechanism combines hierarchical molecular matching (using GNN and attention for multi-level comparisons) and task-level matching (using meta-learning for task adaptation). A detailed explanation and algorithm are now in the appendix A.

  • Q4 Minor issues: We’ve corrected Figure 1 (HF and NH₃) and updated the FS-Mol description in the paper.

• For Reviewer E3W3's concerns:

  • Q1 Outdated baselines: We have clarified the choice of standard baselines and updated the related work to include recent studies like UniCorn and HiPM.

  • Q2 More Pretraining models: We have incorporated the advanced GraphMAE pretraining model, achieving significant performance improvements.

  • Q3 Novelty concerns: We have emphasized our task-matching mechanism as a unique contribution to few-shot learning in molecular representation.

  • Q4 Generalization and dataset imbalance: We have analyzed the extreme class imbalance in the MUV dataset and demonstrated generalization through scaffold-based splits.

  • Q5 Minor issues: We have corrected the query molecule in Figure 2 and clarified the meaning of numerical values in Table 1.

Warm regards,
The Authors of Submission 45

In this submission, the authors propose an effective method for few-shot drug discovery. The proposed method leverages the multi-scale representation of molecules and connects it with implicit task-level generalization, which provides the model performance in various datasets. The proposed method fully leverages the hierarchical structure of molecules and thus makes the representations at different levels work for various tasks, which is reasonable and effective. All the reviewers appreciate this work and are satisfied with the authors' feedback in the rebuttal phase. Taking all these into account, I tend to accept this work.

Accept (Spotlight)